Add set_var to interpreter.py (#263)

sgl-project · Mar 7, 2024 · 30d67b2 · 30d67b2
1 parent b0b722e
commit 30d67b2
Showing 1 changed file with 9 additions and 0 deletions.
diff --git a/python/sglang/lang/interpreter.py b/python/sglang/lang/interpreter.py
@@ -245,6 +245,9 @@ def get_var(self, name):
             self.variable_event[name].wait()
         return self.variables[name]
 
+    def set_var(self, name, value):
+        self.variables[name] = value
+
     def get_meta_info(self, name):
         if name in self.variable_event:
             self.variable_event[name].wait()
@@ -726,6 +729,9 @@ async def text_async_iter(
     def get_var(self, name):
         return self.stream_executor.get_var(name)
 
+    def set_var(self, name, value):
+        return self.stream_executor.set_var(name, value)
+
     def get_meta_info(self, name):
         return self.stream_executor.get_meta_info(name)
 
@@ -736,6 +742,9 @@ def __iadd__(self, other):
     def __getitem__(self, name):
         return self.get_var(name)
 
+    def __setitem__(self, name, value):
+        self.set_var(name, value)
+
     def __del__(self):
         self.stream_executor.end()