@@ -64,6 +64,7 @@ def __init__(
6464 self ._response = {}
6565 self ._message = None
6666 self ._content_block = {}
67+ self ._tool_json_input_buf = ""
6768 self ._record_message = False
6869 self ._ended = False
6970
@@ -89,28 +90,40 @@ def _process_event(self, event):
8990 # {'contentBlockStart': {'start': {'toolUse': {'toolUseId': 'id', 'name': 'func_name'}}, 'contentBlockIndex': 1}}
9091 start = event ["contentBlockStart" ].get ("start" , {})
9192 if "toolUse" in start :
92- tool_use = _decode_tool_use (start ["toolUse" ])
93- self ._content_block = {"toolUse" : tool_use }
93+ self ._content_block = {"toolUse" : start ["toolUse" ]}
9494 return
9595
9696 if "contentBlockDelta" in event :
9797 # {'contentBlockDelta': {'delta': {'text': "Hello"}, 'contentBlockIndex': 0}}
9898 # {'contentBlockDelta': {'delta': {'toolUse': {'input': '{"location":"Seattle"}'}}, 'contentBlockIndex': 1}}
99+ # {'contentBlockDelta': {'delta': {'toolUse': {'input': 'a", "Yok'}}, 'contentBlockIndex': 1}}
99100 if self ._record_message :
100101 delta = event ["contentBlockDelta" ].get ("delta" , {})
101102 if "text" in delta :
102103 self ._content_block .setdefault ("text" , "" )
103104 self ._content_block ["text" ] += delta ["text" ]
104105 elif "toolUse" in delta :
105- tool_use = _decode_tool_use (delta ["toolUse" ])
106- self ._content_block ["toolUse" ].update (tool_use )
106+ if (
107+ input_buf := delta ["toolUse" ].get ("input" )
108+ ) is not None :
109+ self ._tool_json_input_buf += input_buf
107110 return
108111
109112 if "contentBlockStop" in event :
110113 # {'contentBlockStop': {'contentBlockIndex': 0}}
111114 if self ._record_message :
115+ if self ._tool_json_input_buf :
116+ try :
117+ self ._content_block ["toolUse" ]["input" ] = json .loads (
118+ self ._tool_json_input_buf
119+ )
120+ except json .DecodeError :
121+ self ._content_block ["toolUse" ]["input" ] = (
122+ self ._tool_json_input_buf
123+ )
112124 self ._message ["content" ].append (self ._content_block )
113125 self ._content_block = {}
126+ self ._tool_json_input_buf = ""
114127 return
115128
116129 if "messageStop" in event :
0 commit comments