From b8e5684139c62f5774d62165a16c46fc2950ea6e Mon Sep 17 00:00:00 2001 From: bitnom <14287229+bitnom@users.noreply.github.com> Date: Tue, 2 Jan 2024 04:02:43 -0500 Subject: [PATCH 1/4] Handle streamed function calls --- autogen/oai/client.py | 36 ++++++++++++++++++++++++++++-------- 1 file changed, 28 insertions(+), 8 deletions(-) diff --git a/autogen/oai/client.py b/autogen/oai/client.py index 4b548d961ee1..0336e6bc34f0 100644 --- a/autogen/oai/client.py +++ b/autogen/oai/client.py @@ -287,9 +287,8 @@ def yes_or_no_filter(context, response): def _completions_create(self, client, params): completions = client.chat.completions if "messages" in params else client.completions - # If streaming is enabled, has messages, and does not have functions, then - # iterate over the chunks of the response - if params.get("stream", False) and "messages" in params and "functions" not in params: + # If streaming is enabled and has messages, then iterate over the chunks of the response. + if params.get("stream", False) and "messages" in params: response_contents = [""] * params.get("n", 1) finish_reasons = [""] * params.get("n", 1) completion_tokens = 0 @@ -297,12 +296,35 @@ def _completions_create(self, client, params): # Set the terminal text color to green print("\033[32m", end="") + # Prepare for potential function call + full_function_call = None # Send the chat completion request to OpenAI's API and process the response in chunks for chunk in completions.create(**params): if chunk.choices: for choice in chunk.choices: content = choice.delta.content + function_call_chunk = choice.delta.function_call finish_reasons[choice.index] = choice.finish_reason + + # Handle function call + if function_call_chunk: + if hasattr(function_call_chunk, 'name') and function_call_chunk.name: + if full_function_call is None: + full_function_call = {"name": "", "arguments": ""} + full_function_call["name"] += function_call_chunk.name + completion_tokens += 1 + if hasattr(function_call_chunk, 'arguments') and function_call_chunk.arguments: + full_function_call["arguments"] += function_call_chunk.arguments + completion_tokens += 1 + if choice.finish_reason == "function_call": + # Need something here? I don't think so. + pass + if not content: + continue + sys.stdout.write(content) + sys.stdout.flush() + # End handle function call + # If content is present, print it to the terminal and update response variables if content is not None: print(content, end="", flush=True) @@ -336,7 +358,7 @@ def _completions_create(self, client, params): index=i, finish_reason=finish_reasons[i], message=ChatCompletionMessage( - role="assistant", content=response_contents[i], function_call=None + role="assistant", content=response_contents[i], function_call=full_function_call ), logprobs=None, ) @@ -346,16 +368,14 @@ def _completions_create(self, client, params): index=i, finish_reason=finish_reasons[i], message=ChatCompletionMessage( - role="assistant", content=response_contents[i], function_call=None + role="assistant", content=response_contents[i], function_call=full_function_call ), ) response.choices.append(choice) else: - # If streaming is not enabled or using functions, send a regular chat completion request - # Functions are not supported, so ensure streaming is disabled + # If streaming is not enabled, send a regular chat completion request params = params.copy() - params["stream"] = False response = completions.create(**params) return response From 5bd9d77b23368fadeb00a1383acca21e4bbfb38f Mon Sep 17 00:00:00 2001 From: bitnom <14287229+bitnom@users.noreply.github.com> Date: Tue, 2 Jan 2024 13:51:46 -0500 Subject: [PATCH 2/4] apply black formatting --- autogen/oai/client.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/autogen/oai/client.py b/autogen/oai/client.py index 0336e6bc34f0..bec0361226e2 100644 --- a/autogen/oai/client.py +++ b/autogen/oai/client.py @@ -308,12 +308,12 @@ def _completions_create(self, client, params): # Handle function call if function_call_chunk: - if hasattr(function_call_chunk, 'name') and function_call_chunk.name: + if hasattr(function_call_chunk, "name") and function_call_chunk.name: if full_function_call is None: full_function_call = {"name": "", "arguments": ""} full_function_call["name"] += function_call_chunk.name completion_tokens += 1 - if hasattr(function_call_chunk, 'arguments') and function_call_chunk.arguments: + if hasattr(function_call_chunk, "arguments") and function_call_chunk.arguments: full_function_call["arguments"] += function_call_chunk.arguments completion_tokens += 1 if choice.finish_reason == "function_call": From 9844c69cba5076076a4c485af2b5cd5211e7786f Mon Sep 17 00:00:00 2001 From: bitnom <14287229+bitnom@users.noreply.github.com> Date: Tue, 2 Jan 2024 15:16:15 -0500 Subject: [PATCH 3/4] rm unnecessary stdout print --- autogen/oai/client.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/autogen/oai/client.py b/autogen/oai/client.py index bec0361226e2..1d2087401547 100644 --- a/autogen/oai/client.py +++ b/autogen/oai/client.py @@ -321,8 +321,6 @@ def _completions_create(self, client, params): pass if not content: continue - sys.stdout.write(content) - sys.stdout.flush() # End handle function call # If content is present, print it to the terminal and update response variables From 8c284b16abe93b5dd7f6927cfe1a10e6a14dbe2c Mon Sep 17 00:00:00 2001 From: Davor Runje Date: Wed, 3 Jan 2024 09:08:22 +0000 Subject: [PATCH 4/4] bug fix --- autogen/oai/client.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/autogen/oai/client.py b/autogen/oai/client.py index 1d2087401547..ce36f8e10180 100644 --- a/autogen/oai/client.py +++ b/autogen/oai/client.py @@ -374,7 +374,9 @@ def _completions_create(self, client, params): else: # If streaming is not enabled, send a regular chat completion request params = params.copy() + params["stream"] = False response = completions.create(**params) + return response def _update_usage_summary(self, response: ChatCompletion | Completion, use_cache: bool) -> None: