From 916917a9bf4927eea8abd0e85ad12e3203121bdc Mon Sep 17 00:00:00 2001 From: 苏向夜 Date: Sat, 24 Feb 2024 21:52:46 +0800 Subject: feat(input): add output method for input --- src/infini/input.py | 24 ++++++++++++++++++++++-- 1 file changed, 22 insertions(+), 2 deletions(-) diff --git a/src/infini/input.py b/src/infini/input.py index 682e1037..30683a4b 100644 --- a/src/infini/input.py +++ b/src/infini/input.py @@ -1,14 +1,21 @@ -from infini.typing import Dict, Any, Generic, T +from infini.typing import Literal, Optional +from infini.typing import Dict, Any, Generic, T, Optional +from infini.output import Output class Input(Generic[T]): plain_data: T variables: Dict[str, Any] - def __init__(self, plain_data: T, variables: Dict[str, Any] | None = None) -> None: + def __init__( + self, plain_data: T, variables: Optional[Dict[str, Any]] = None + ) -> None: self.plain_data = plain_data self.variables = variables or {} + def get_user_id(self) -> Optional[str]: + return self.variables.get("user_id") + def get_session_id(self) -> str: if session_id := self.variables.get("session_id"): return session_id @@ -19,3 +26,16 @@ class Input(Generic[T]): def get_plain_text(self) -> str: return str(self.plain_data) + + def output( + self, + type: Literal["text", "workflow"], + name: str, + *, + status: int = 0, + block: bool = False, + variables: Dict[str, Any] = {}, + ): + vars = self.variables.copy() + vars.update(variables) + return Output(type, name, status=status, block=block, variables=vars) -- cgit v1.2.3-70-g09d2