Abhaykoul commited on
Commit
019c8ae
1 Parent(s): 3fc5c4e

Create stream.py

Browse files
Files changed (1) hide show
  1. stream.py +22 -0
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
+