Merge pull request #43 from atomiechen/fix_tool_call

Fix tool call

src/handyllm/hprompt.py

@@ -731,6 +731,10 @@ def _stream_with_client(
                     run_config.on_chunk = cast(SyncHandlerChat, run_config.on_chunk)
                     run_config.on_chunk(*ret)
                 yield chat_chunk
+            ret = producer.send(None)  # signal the end of the stream
+            if run_config.on_chunk and ret:
+                run_config.on_chunk = cast(SyncHandlerChat, run_config.on_chunk)
+                run_config.on_chunk(*ret)
             producer.close()
 
     @classmethod
@@ -756,6 +760,13 @@ async def _astream_with_client(
                         run_config.on_chunk = cast(SyncHandlerChat, run_config.on_chunk)
                         run_config.on_chunk(*ret)
                 yield chat_chunk
+            ret = producer.send(None)  # signal the end of the stream
+            if run_config.on_chunk and ret:
+                if inspect.iscoroutinefunction(run_config.on_chunk):
+                    await run_config.on_chunk(*ret)
+                else:
+                    run_config.on_chunk = cast(SyncHandlerChat, run_config.on_chunk)
+                    run_config.on_chunk(*ret)
             producer.close()
 
     @classmethod

src/handyllm/utils.py

@@ -21,6 +21,7 @@
 ]
 
 import base64
+import copy
 from pathlib import Path
 from typing import (
     IO,
@@ -69,14 +70,16 @@ def download_binary(download_url, file_path=None, dir="."):
 
 def trans_stream_chat(
     consumer: Generator[YieldType, ShortChatChunk, None],
-) -> Generator[Optional[YieldType], ChatChunk, None]:
+) -> Generator[Optional[YieldType], Optional[ChatChunk], None]:
     next(consumer)  # prime the generator
     role = ""
     tool_call = ToolCallDelta()
     ret = None
     try:
         while True:
             data = yield ret
+            if data is None:
+                break
             ret = None
             try:
                 message = data["choices"][0]["delta"]
@@ -93,19 +96,24 @@ def trans_stream_chat(
                                 "arguments"
                             ]
                         else:
-                            # this is a new tool call, yield the previous one
-                            ret = consumer.send((role, content, tool_call))
+                            if tool_call:
+                                # this is a new tool call, yield the previous one
+                                ret = consumer.send((role, content, tool_call))
                             # reset the tool call
-                            tool_call = ToolCallDelta(chunk)
+                            tool_call = copy.deepcopy(chunk)
                 elif content:
                     ret = consumer.send((role, content, tool_call))
             except (KeyError, IndexError):
                 pass
-    except GeneratorExit:
         if tool_call:
             # yield the last tool call
             ret = consumer.send((role, None, tool_call))
+            yield ret
+        else:
+            yield None
         consumer.close()
+    except GeneratorExit:
+        pass
 
 
 def echo_consumer():
@@ -123,6 +131,9 @@ def stream_chat_all(
         ret = producer.send(data)
         if ret is not None:
             yield ret
+    ret = producer.send(None)  # signal the end of the stream
+    if ret is not None:
+        yield ret
     producer.close()
 
 
@@ -154,6 +165,9 @@ async def astream_chat_all(
         ret = producer.send(data)
         if ret is not None:
             yield ret
+    ret = producer.send(None)  # signal the end of the stream
+    if ret is not None:
+        yield ret
     producer.close()
 
 

tests/assets/chat_tool.hprompt

@@ -0,0 +1,37 @@
+---
+meta:
+  var_map_path: var_map.txt
+  output_path: tmp-out/%Y-%m-%d/result.%H-%M-%S.hprompt
+  output_evaled_prompt_path: tmp-evaled/%Y-%m-%d/result.%H-%M-%S.hprompt
+model: gpt-4o
+temperature: 0.2
+tools: [
+    {
+      "type": "function",
+      "function": {
+        "name": "get_current_weather",
+        "description": "Get the current weather in a given location",
+        "parameters": {
+          "type": "object",
+          "properties": {
+            "location": {
+              "type": "string",
+              "description": "The city and state, e.g. San Francisco, CA"
+            },
+            "unit": {
+              "type": "string",
+              "enum": ["celsius", "fahrenheit"]
+            }
+          },
+          "required": ["location", "unit"]
+        }
+      }
+    }
+  ]
+---
+
+$system$
+You are a helpful assistant.
+
+$user$
+What's the weather like in San Francisco and New York?