EngineX-Hygon/sglang
5
0
Fork 0
Code Issues Pull Requests Actions 7 Projects Releases Wiki Activity

Add option to return metadata in async streaming (#18)

Browse Source
This commit is contained in:
Christopher Chou
2024-01-17 18:15:02 -08:00
committed by GitHub
parent 8024fc5eec
commit c0454b323c
1 changed files with 6 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files

10
python/sglang/lang/interpreter.py
View File

@@ -365,11 +365,10 @@ class StreamExecutor:
for comp, meta_info in generator: for comp, meta_info in generator:
self.text_ += comp self.text_ += comp
self.variables[name] += comp self.variables[name] += comp
self.meta_info[name] = meta_info
self.stream_var_event[name].set() self.stream_var_event[name].set()
self.stream_text_event.set() self.stream_text_event.set()
self.meta_info[name] = meta_info
self.variable_event[name].set() self.variable_event[name].set()
self.stream_var_event[name].set() self.stream_var_event[name].set()
@@ -583,7 +582,7 @@ class ProgramState:
else: else:
yield self.get_var(name) yield self.get_var(name)
async def text_async_iter(self, var_name=None): async def text_async_iter(self, var_name=None, return_meta_data=False):
loop = asyncio.get_running_loop() loop = asyncio.get_running_loop()
if self.stream_executor.stream: if self.stream_executor.stream:
@@ -607,7 +606,10 @@ class ProgramState:
out = str(self.stream_executor.variables[var_name][prev:]) out = str(self.stream_executor.variables[var_name][prev:])
prev += len(out) prev += len(out)
if out: if out:
yield out if return_meta_data:
yield out, self.stream_executor.meta_info[var_name]
else:
yield out
if self.stream_executor.variable_event[var_name].is_set(): if self.stream_executor.variable_event[var_name].is_set():
break break
else: else: