Merge remote-tracking branch 'refs/remotes/origin/master'

Author: iamlemec

oneping/native/openai.py

@@ -28,29 +28,29 @@ def make_client(api_key=None, async_client=False):
     client_class = openai.AsyncOpenAI if async_client else openai.OpenAI
     return client_class(api_key=api_key)
 
-def reply(query, image=None, history=None, prefill=None, prediction=None, system=DEFAULT_SYSTEM, api_key=None, model=OPENAI_MODEL, **kwargs):
+def reply(query, image=None, history=None, prefill=None, prediction=None, system=DEFAULT_SYSTEM, api_key=None, model=OPENAI_MODEL, max_tokens=None, **kwargs):
     client = make_client(api_key=api_key)
     payload = make_payload(query, image=image, prediction=prediction, system=system, history=history)
-    response = client.chat.completions.create(model=model, **payload, **kwargs)
+    response = client.chat.completions.create(model=model, max_completion_tokens=max_tokens, **payload, **kwargs)
     return response_openai_native(response)
 
-async def reply_async(query, image=None, history=None, prefill=None, prediction=None, system=DEFAULT_SYSTEM, api_key=None, model=OPENAI_MODEL, **kwargs):
+async def reply_async(query, image=None, history=None, prefill=None, prediction=None, system=DEFAULT_SYSTEM, api_key=None, model=OPENAI_MODEL, max_tokens=None, **kwargs):
     client = make_client(api_key=api_key, async_client=True)
     payload = make_payload(query, image=image, prediction=prediction, system=system, history=history)
-    response = await client.chat.completions.create(model=model, **payload, **kwargs)
+    response = await client.chat.completions.create(model=model, max_completion_tokens=max_tokens, **payload, **kwargs)
     return response_openai_native(response)
 
-def stream(query, image=None, history=None, prefill=None, prediction=None, system=DEFAULT_SYSTEM, api_key=None, model=OPENAI_MODEL, **kwargs):
+def stream(query, image=None, history=None, prefill=None, prediction=None, system=DEFAULT_SYSTEM, api_key=None, model=OPENAI_MODEL, max_tokens=None, **kwargs):
     client = make_client(api_key=api_key)
     payload = make_payload(query, image=image, prediction=prediction, system=system, history=history)
-    response = client.chat.completions.create(model=model, stream=True, **payload, **kwargs)
+    response = client.chat.completions.create(model=model, stream=True, max_completion_tokens=max_tokens, **payload, **kwargs)
     for chunk in response:
         yield stream_openai_native(chunk)
 
-async def stream_async(query, image=None, history=None, prefill=None, prediction=None, system=DEFAULT_SYSTEM, api_key=None, model=OPENAI_MODEL, **kwargs):
+async def stream_async(query, image=None, history=None, prefill=None, prediction=None, system=DEFAULT_SYSTEM, api_key=None, model=OPENAI_MODEL, max_tokens=None, **kwargs):
     client = make_client(api_key=api_key, async_client=True)
     payload = make_payload(query, image=image, prediction=prediction, system=system, history=history)
-    response = await client.chat.completions.create(model=model, stream=True, **payload, **kwargs)
+    response = await client.chat.completions.create(model=model, stream=True, max_completion_tokens=max_tokens, **payload, **kwargs)
     async for chunk in response:
         yield stream_openai_native(chunk)