router pushes sse events; better streaming extraction

Author: iamlemec

oneping/curl.py

@@ -12,18 +12,15 @@
 ##
 
 def prepare_url(prov, url=None, host=None, port=None):
-    if host is None:
-        host = 'localhost'
-    if port is None:
-        port = 8000
-    if url is None:
-        url = prov['url'].format(host=host, port=port)
-    return url
+    host = prov.get('host') if host is None else host
+    port = prov.get('port') if port is None else port
+    url = prov.get('url') if url is None else url
+    return url.format(host=host, port=port)
 
 def prepare_auth(prov, api_key=None):
     if (auth_func := prov.get('authorize')) is not None:
-        if api_key is None and (api_key := os.environ.get(key_env := prov['api_key_env'])) is None:
-            raise Exception('Cannot find API key in {key_env}')
+        if (api_key := os.environ.get(prov['api_key_env'])) is None:
+            raise Exception('Cannot find API key in {api_key_env}')
         headers_auth = auth_func(api_key)
     else:
         headers_auth = {}
@@ -36,7 +33,7 @@ def prepare_model(prov, model=None):
 
 def prepare_request(
     query, provider='local', system=None, prefill=None, prediction=None, history=None,
-    url=None, port=None, api_key=None, model=None, max_tokens=DEFAULT_MAX_TOKENS, **kwargs
+    url=None, host=None, port=None, api_key=None, model=None, max_tokens=DEFAULT_MAX_TOKENS, **kwargs
 ):
     # external provider
     prov = get_provider(provider)
@@ -45,7 +42,7 @@ def prepare_request(
     max_tokens_name = prov.get('max_tokens_name', 'max_tokens')
 
     # get full url
-    url = prepare_url(prov, url=url, port=port)
+    url = prepare_url(prov, url=url, host=host, port=port)
 
     # get authorization headers
     headers_auth = prepare_auth(prov, api_key=api_key)
@@ -127,9 +124,9 @@ async def reply_async(query, provider='local', history=None, prefill=None, **kwa
 ## stream requests
 ##
 
-def parse_stream_data(chunk):
-    if chunk.startswith(b'data: '):
-        text = chunk[6:]
+def parse_stream_data(line):
+    if line.startswith(b'data: '):
+        text = line[6:]
         if text != b'[DONE]' and len(text) > 0:
             return text
 
@@ -172,7 +169,9 @@ def stream(query, provider='local', history=None, prefill=None, **kwargs):
         for line in response.iter_lines():
             if (data := parse_stream_data(line)) is not None:
                 parsed = json.loads(data)
-                yield extractor(parsed)
+                text = extractor(parsed)
+                if text is not None:
+                    yield text
 
 async def stream_async(query, provider='local', history=None, prefill=None, **kwargs):
     # get provider
@@ -206,7 +205,9 @@ async def stream_async(query, provider='local', history=None, prefill=None, **kw
             async for line in lines:
                 if (data := parse_stream_data(line)) is not None:
                     parsed = json.loads(data)
-                    yield extractor(parsed)
+                    text = extractor(parsed)
+                    if text is not None:
+                        yield text
 
 ##
 ## embeddings

oneping/providers.py

@@ -21,6 +21,7 @@
 ANTHROPIC_MODEL = 'claude-3-5-sonnet-latest'
 FIREWORKS_MODEL = 'accounts/fireworks/models/llama-v3p1-70b-instruct'
 GROQ_MODEL = 'llama-3.1-70b-versatile'
+DEEPSEEK_MODEL = 'deepseek-chat'
 
 ##
 ## environment key names
@@ -30,6 +31,7 @@
 ANTHROPIC_KEYENV = 'ANTHROPIC_API_KEY'
 FIREWORKS_KEYENV = 'FIREWORKS_API_KEY'
 GROQ_KEYENV = 'GROQ_API_KEY'
+DEEPSEEK_KEYENV = 'DEEPSEEK_API_KEY'
 AZURE_KEYENV = 'AZURE_OPENAI_API_KEY'
 
 ##
@@ -114,17 +116,23 @@ def response_anthropic(reply):
     content = reply['content'][0]
     return content['text']
 
-def stream_oneping(chunk):
-    return chunk
+##
+## stream handlers
+##
 
 def stream_openai(chunk):
     return chunk['choices'][0]['delta'].get('content', '')
 
 def stream_anthropic(chunk):
     if chunk['type'] == 'content_block_delta':
         return chunk['delta']['text']
-    else:
-        return ''
+
+def stream_oneping(chunk):
+    return chunk
+
+##
+## native handlers
+##
 
 def response_openai_native(reply):
     return reply.choices[0].message.content
@@ -145,6 +153,10 @@ def stream_anthropic_native(chunk):
     else:
         return ''
 
+##
+## other modal handlers
+##
+
 def embed_openai(reply):
     return reply['data'][0]['embedding']
 
@@ -156,7 +168,6 @@ def transcribe_openai(audio):
 ##
 
 DEFAULT_PROVIDER = {
-    'authorize': authorize_openai,
     'payload': payload_openai,
     'response': response_openai,
     'stream': stream_openai,
@@ -167,18 +178,21 @@ def transcribe_openai(audio):
 LLM_PROVIDERS = {
     'local': {
         'url': 'http://{host}:{port}/v1/chat/completions',
-        'authorize': None,
+        'host': 'localhost',
+        'port': 8000,
     },
     'oneping': {
         'url': 'http://{host}:{port}/chat',
         'host': 'localhost',
         'port': 5000,
+        'authorize': None,
         'payload': payload_oneping,
         'response': response_oneping,
         'stream': stream_oneping,
     },
     'openai': {
         'url': 'https://api.openai.com/v1/chat/completions',
+        'authorize': authorize_openai,
         'max_tokens_name': 'max_completion_tokens',
         'api_key_env': OPENAI_KEYENV,
         'model': OPENAI_MODEL,
@@ -195,14 +209,22 @@ def transcribe_openai(audio):
     },
     'fireworks': {
         'url': 'https://api.fireworks.ai/inference/v1/chat/completions',
+        'authorize': authorize_openai,
         'api_key_env': FIREWORKS_KEYENV,
         'model': FIREWORKS_MODEL,
     },
     'groq': {
         'url': 'https://api.groq.com/openai/v1/chat/completions',
+        'authorize': authorize_openai,
         'api_key_env': GROQ_KEYENV,
         'model': GROQ_MODEL,
     },
+    'deepseek': {
+        'url': 'https://api.deepseek.com/chat/completions',
+        'authorize': authorize_openai,
+        'api_key_env': DEEPSEEK_KEYENV,
+        'model': DEEPSEEK_MODEL,
+    },
 }
 
 def get_provider(provider):

oneping/server.py

@@ -1,5 +1,6 @@
 # llm servers
 
+import json
 import subprocess
 from itertools import chain
 
@@ -26,6 +27,12 @@ def patch_payload(data):
             data['provider'] = model
     return data
 
+def generate_sse(stream):
+    for chunk in stream:
+        data = json.dumps(chunk)
+        yield f'data: {data}\n\n'
+    yield 'data: [DONE]\n\n'
+
 def start_router(host='127.0.0.1', port=5000, allow_origins=DEFAULT_ALLOW_ORIGINS, **kwargs):
     import uvicorn
     from fastapi import FastAPI
@@ -61,11 +68,11 @@ async def chat(genreq: GenerateRequest):
         data = genreq.model_dump(exclude_none=True)
         patch = patch_payload(data)
         if patch.get('stream', False):
-            ret = stream_api(**kwargs, **patch)
-            return StreamingResponse(ret, media_type='text/plain')
+            stream = stream_api(**kwargs, **patch)
+            sse = generate_sse(stream)
+            return StreamingResponse(sse, media_type='text/event-stream')
         else:
-            ret = reply_api(**kwargs, **patch)
-            text = ret[1] if type(ret) is tuple else ret
-            return PlainTextResponse(text)
+            reply = reply_api(**kwargs, **patch)
+            return PlainTextResponse(reply)
 
     uvicorn.run(app, host=host, port=port)