Update.

Browse Source
This commit is contained in:
retoor 2025-04-26 12:17:24 +02:00
parent f1c4553038
commit 598cf35886
2 changed files with 33 additions and 20 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

16
client.py
View File

@ -14,35 +14,29 @@ async def websocket_client(url: str, ollama_url: str) -> None:
async with aiohttp.ClientSession() as session: async with aiohttp.ClientSession() as session:
try: try:
async with session.ws_connect(f'{url}/publish') as ws: async with session.ws_connect(f'{url}/publish') as ws:
logging.info("Fetching models.")
async with session.get(f'{ollama_url}/api/tags') as response: async with session.get(f'{ollama_url}/api/tags') as response:
if response.status != 200: if response.status != 200:
logging.error(f"Failed to fetch models: {response.status}") logging.error(f"Failed to fetch models: {response.status}")
return return
models = await response.json() models = await response.json()
await ws.send_json(models) await ws.send_json(models)
logging.info("Published models to uberlama.")
async for msg in ws: async for msg in ws:
if msg.type == aiohttp.WSMsgType.TEXT: if msg.type == aiohttp.WSMsgType.TEXT:
data = msg.json() data = msg.json()
logging.info(f"Received data: {data}") logging.info(f"Received data: {data}.")
request_id = data['request_id'] request_id = data['request_id']
api_url = urlunparse(urlparse(ollama_url)._replace(path=data['path'])) api_url = urlunparse(urlparse(ollama_url)._replace(path=data['path']))
async with session.post(api_url, json=data['data']) as response: async with session.post(api_url, json=data['data']) as response:
print(response)
if response.status != 200: if response.status != 200:
logging.error(f"Failed to post data: {response.status}") logging.error(f"Failed to post data: {response.status}.")
continue continue
logging.info(f"Streaming response.") logging.info(f"Streaming response.")
async for msg in response.content: async for msg in response.content:
#first_index = msg.find(b"{") print(msg.decode())
#msg = msg[first_index:]
#last_index = msg.rfind(b"}")
#msg = msg[:last_index+1]
#if not msg:
# continue
#msg = json.loads(msg.decode('utf-8'))
await ws.send_json(dict( await ws.send_json(dict(
request_id=request_id, request_id=request_id,
data=msg.decode() data=msg.decode()

37
server.py
View File

@ -31,18 +31,37 @@ class OllamaServer:
if chunk: if chunk:
yield chunk yield chunk
if not chunk: if not chunk:
yield '\n' yield ''
#yield '\n'
print("CHUNK:", chunk) print("CHUNK:", chunk)
try:
obj = json.loads(chunk)
if obj.get('done'):
break
except:
pass
try:
if '"finish_reason":"stop"' in chunk:
break
except:
pass
try:
if 'data: [DONE]' in chunk:
break
except:
pass
#try: #try:
#yield json.loads(chunk) #yield json.loads(chunk)
#except: #except:
# yield chunk # yield chunk
if not 'done' in chunk: #if not 'done' in chunk:
break # break
if 'stop' in chunk: #if 'stop' in chunk:
break # break
if chunk['done']: #if chunk.get('done'):
break # break
async def serve(self): async def serve(self):
@ -145,10 +164,10 @@ async def http_handler(request):
except ValueError: except ValueError:
return web.Response(status=400) return web.Response(status=400)
# application/x-ndjson text/event-stream # application/x-ndjson text/event-stream
if data['stream']: if data.get('stream'):
resp = web.StreamResponse(headers={'Content-Type': 'text/event-stream', 'Transfer-Encoding': 'chunked'}) resp = web.StreamResponse(headers={'Content-Type': 'text/event-stream', 'Transfer-Encoding': 'chunked'})
else: else:
resp = web.StreamResponse(headers={'Content-Type': 'application/json', 'Transfer-Encoding': 'chunked'}) resp = web.StreamResponse(headers={'Content-Type': 'application/json', 'Transfer-Encoding':'chunked'})
await resp.prepare(request) await resp.prepare(request)
import json import json
try: try: