From c0454b323c87c78f650074659c2db92b0fd2524d Mon Sep 17 00:00:00 2001
From: Christopher Chou <49086305+BabyChouSr@users.noreply.github.com>
Date: Wed, 17 Jan 2024 18:15:02 -0800
Subject: [PATCH] Add option to return metadata in async streaming (#18)

---
 python/sglang/lang/interpreter.py | 10 ++++++----
 1 file changed, 6 insertions(+), 4 deletions(-)

diff --git a/python/sglang/lang/interpreter.py b/python/sglang/lang/interpreter.py
index 302ec96c5..150fd04fc 100644
--- a/python/sglang/lang/interpreter.py
+++ b/python/sglang/lang/interpreter.py
@@ -365,11 +365,10 @@ class StreamExecutor:
             for comp, meta_info in generator:
                 self.text_ += comp
                 self.variables[name] += comp
+                self.meta_info[name] = meta_info
                 self.stream_var_event[name].set()
                 self.stream_text_event.set()
 
-            self.meta_info[name] = meta_info
-
             self.variable_event[name].set()
             self.stream_var_event[name].set()
 
@@ -583,7 +582,7 @@ class ProgramState:
             else:
                 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()
 
         if self.stream_executor.stream:
@@ -607,7 +606,10 @@ class ProgramState:
                     out = str(self.stream_executor.variables[var_name][prev:])
                     prev += len(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():
                         break
         else: