Create stream.py
Browse files
stream.py
ADDED
@@ -0,0 +1,22 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
import asyncio
|
2 |
+
import json
|
3 |
+
import aiohttp
|
4 |
+
|
5 |
+
async def fastai_stream(user, model="llama3-8b", system="Answer as concisely as possible."):
|
6 |
+
env_type = "tp16405b" if "405b" in model else "tp16"
|
7 |
+
data = {'body': {'messages': [{'role': 'system', 'content': system}, {'role': 'user', 'content': user}], 'stream': True, 'model': model}, 'env_type': env_type}
|
8 |
+
|
9 |
+
async with aiohttp.ClientSession() as session:
|
10 |
+
async with session.post('https://fast.snova.ai/api/completion', headers={'content-type': 'application/json'}, json=data) as response:
|
11 |
+
async for line in response.content:
|
12 |
+
line = line.decode('utf-8').strip()
|
13 |
+
if line.startswith('data:'):
|
14 |
+
try:
|
15 |
+
data = json.loads(line[len('data: '):])
|
16 |
+
content = data.get("choices", [{}])[0].get("delta", {}).get("content", '')
|
17 |
+
if content:
|
18 |
+
yield f"data: {json.dumps({'response': content})}\n\n"
|
19 |
+
except json.JSONDecodeError:
|
20 |
+
if line[len('data: '):] == '[DONE]':
|
21 |
+
break
|
22 |
+
|