feat(injector): supports backend function injector

4 files changed, 44 insertions, 14 deletions

diff --git a/src/infini/core.py b/src/infini/core.py
index fe99f4b3..08b81f15 100644
--- a/src/infini/core.py
+++ b/src/infini/core.py
@@ -40,7 +40,7 @@ class Core:
             if handled_stream.is_empty():
                 return
             outcome = self.generate(handled_stream)
-            for stream in self.intercept(outcome):
+            for stream in self.intercept(handled_stream, outcome):
                 if isinstance(stream, Output):
                     if stream.is_empty():
                         return
@@ -64,5 +64,7 @@ class Core:
     def generate(self, output: Output) -> str:
         return self.generator.output(output, self.injector)
 
-    def intercept(self, output_text: str) -> Generator[Union[str, Output], Any, None]:
-        return self.interceptor.output(output_text)
+    def intercept(
+        self, output: Output, output_text: str
+    ) -> Generator[Union[str, Output], Any, None]:
+        return self.interceptor.output(output, output_text)
diff --git a/src/infini/handler.py b/src/infini/handler.py
index 0d8d1fe0..3d8cd586 100644
--- a/src/infini/handler.py
+++ b/src/infini/handler.py
@@ -1,3 +1,4 @@
+from infini.injector import Injector
 from infini.input import Input
 from infini.output import Output
 from infini.typing import List, Any, RouterType, Callable, Generator, Union
@@ -11,8 +12,16 @@ class Handler:
         if (queue := self.match(input.get_plain_text())).is_empty():
             yield Output.empty()
             return
+        injector = Injector()
+        parameters = {
+            "input": input,
+            "plain_text": input.get_plain_text(),
+            "user_id": input.get_user_id(),
+        }
         while not queue.is_empty():
-            if isinstance(stream := queue.pop()(input), Generator):
+            if isinstance(
+                stream := injector.output(queue.pop(), parameters=parameters), Generator
+            ):
                 for output in stream:
                     yield output
                     if output.block:
@@ -24,7 +33,7 @@ class Handler:
 
     def match(
         self, text: str
-    ) -> EventQueue[Callable[[Input], Union[Output, Generator[Output, Any, None]]]]:
+    ) -> EventQueue[Callable[..., Union[Output, Generator[Output, Any, None]]]]:
         queue = EventQueue()
 
         for handler in self.handlers:
diff --git a/src/infini/interceptor.py b/src/infini/interceptor.py
index bffd3f8e..87255ae9 100644
--- a/src/infini/interceptor.py
+++ b/src/infini/interceptor.py
@@ -1,3 +1,4 @@
+from infini.injector import Injector
 from infini.input import Input
 from infini.output import Output
 from infini.typing import List, Any, RouterType, Callable, Generator, Union
@@ -9,8 +10,17 @@ class Interceptor:
 
     def input(self, input: Input) -> Generator[Union[Output, Input], Any, None]:
         queue = self.match(input.get_plain_text())
+        injector = Injector()
+        parameters = {
+            "input": input,
+            "plain_text": input.get_plain_text(),
+            "user_id": input.get_user_id(),
+        }
         while not queue.is_empty():
-            if isinstance(stream := queue.pop()(input), Generator):
+            if isinstance(
+                stream := injector.output(queue.pop(), parameters=parameters),
+                Generator,
+            ):
                 for outcome in stream:
                     if isinstance(outcome, Input):
                         input = outcome
@@ -30,11 +40,22 @@ class Interceptor:
                 input = stream
         yield input
 
-    def output(self, output_text: str) -> Generator[Union[Output, str], Any, None]:
-        input = Input(output_text)
-        queue = self.match(input.get_plain_text())
+    def output(
+        self, output: Output, output_text: str
+    ) -> Generator[Union[Output, str], Any, None]:
+        input = Input(output_text, variables=output.variables)
+        queue = self.match(output_text)
+        injector = Injector()
+        parameters = {
+            "input": input,
+            "output": output,
+            "plain_text": output_text,
+            "user_id": input.get_user_id(),
+        }
         while not queue.is_empty():
-            if isinstance(stream := queue.pop()(input), Generator):
+            if isinstance(
+                stream := injector.output(queue.pop(), parameters=parameters), Generator
+            ):
                 for outcome in stream:
                     if isinstance(outcome, Input):
                         input = outcome
@@ -57,9 +78,7 @@ class Interceptor:
     def match(
         self, text: str
     ) -> EventQueue[
-        Callable[
-            [Input], Union[Input, Output, Generator[Union[Input, Output], Any, None]]
-        ]
+        Callable[..., Union[Input, Output, Generator[Union[Input, Output], Any, None]]]
     ]:
         queue = EventQueue()
 
diff --git a/src/infini/typing.py b/src/infini/typing.py
index 6bac0099..874e8753 100644
--- a/src/infini/typing.py
+++ b/src/infini/typing.py
@@ -24,4 +24,4 @@ OutputGenerator = Generator["output.Output", Any, None]
 class RouterType(TypedDict):
     priority: int
     router: router.Router
-    handler: Callable[["input.Input"], Union[Stream, OutputGenerator]]
+    handler: Callable[..., Union[Stream, OutputGenerator]]