Feat: add mechanism to check cancellation in Agent (#10766)

### What problem does this PR solve?

Add mechanism to check cancellation in Agent.

### Type of change

- [x] New Feature (non-breaking change which adds functionality)
This commit is contained in:
Yongteng Lei
2025-11-11 17:36:48 +08:00
committed by GitHub
parent d81e4095de
commit 9213568692
36 changed files with 495 additions and 20 deletions

View File

@ -89,6 +89,9 @@ class Message(ComponentBase):
all_content = ""
cache = {}
for r in re.finditer(self.variable_ref_patt, rand_cnt, flags=re.DOTALL):
if self.check_if_canceled("Message streaming"):
return
all_content += rand_cnt[s: r.start()]
yield rand_cnt[s: r.start()]
s = r.end()
@ -104,6 +107,9 @@ class Message(ComponentBase):
if isinstance(v, partial):
cnt = ""
for t in v():
if self.check_if_canceled("Message streaming"):
return
all_content += t
cnt += t
yield t
@ -111,7 +117,7 @@ class Message(ComponentBase):
continue
elif not isinstance(v, str):
try:
v = json.dumps(v, ensure_ascii=False, indent=2)
v = json.dumps(v, ensure_ascii=False)
except Exception:
v = str(v)
yield v
@ -120,6 +126,9 @@ class Message(ComponentBase):
cache[exp] = v
if s < len(rand_cnt):
if self.check_if_canceled("Message streaming"):
return
all_content += rand_cnt[s: ]
yield rand_cnt[s: ]
@ -133,6 +142,9 @@ class Message(ComponentBase):
@timeout(int(os.environ.get("COMPONENT_EXEC_TIMEOUT", 10*60)))
def _invoke(self, **kwargs):
if self.check_if_canceled("Message processing"):
return
rand_cnt = random.choice(self._param.content)
if self._param.stream and not self._is_jinjia2(rand_cnt):
self.set_output("content", partial(self._stream, rand_cnt))
@ -145,6 +157,9 @@ class Message(ComponentBase):
except Exception:
pass
if self.check_if_canceled("Message processing"):
return
for n, v in kwargs.items():
content = re.sub(n, v, content)