FilesExpand file tree

 main

/

test_stream.py

Copy path

Blame

More file actions

Blame

More file actions

Latest commit

 

History

History

History

55 lines (40 loc) · 1.28 KB

 main

/

test_stream.py

Copy path

Top

File metadata and controls

• Code
• Blame

55 lines (40 loc) · 1.28 KB

Raw

Copy raw file

Download raw file

Open symbols panel

Edit and raw actions

1

2

3

4

5

6

7

8

9

10

11

12

13

14

15

16

17

18

19

20

21

22

23

24

25

26

27

28

29

30

31

32

33

34

35

36

37

38

39

40

41

42

43

44

45

46

47

48

49

50

51

52

53

54

55

import pytest

import replicate

@pytest.mark.asyncio

@pytest.mark.parametrize("async_flag", [True, False])

async def test_stream(async_flag, record_mode):

if record_mode == "none":

return

version = "02e509c789964a7ea8736978a43525956ef40397be9033abf9fd2badfe68c9e3"

input = {

"prompt": "Please write a haiku about llamas.",

}

events = []

if async_flag:

async for event in await replicate.async_stream(

f"meta/llama-2-70b-chat:{version}",

input=input,

):

events.append(event)

else:

for event in replicate.stream(

f"meta/llama-2-70b-chat:{version}",

input=input,

):

events.append(event)

assert len(events) > 0

assert events[0].event == "output"

@pytest.mark.asyncio

async def test_stream_prediction(record_mode):

if record_mode == "none":

return

version = "02e509c789964a7ea8736978a43525956ef40397be9033abf9fd2badfe68c9e3"

input = {

"prompt": "Please write a haiku about llamas.",

}

prediction = replicate.predictions.create(version=version, input=input)

events = []

for event in prediction.stream():

events.append(event)

assert len(events) > 0

assert events[0].event == "output"