[AUR-395] Adopt Example1 disclaimer pipeline (#42)

* Adopt Example1 disclaimer pipeline * Update Document class * Add composite components * Modify Extractor behaviours
2023-10-10 15:42:48 +07:00
parent 79cc60e6a2
commit 84f1fa8cbd
12 changed files with 654 additions and 37 deletions
--- a/knowledgehub/composite/linear.py
+++ b/knowledgehub/composite/linear.py
@@ -0,0 +1,153 @@
+from typing import Any, Callable, Optional, Union
+
+from kotaemon.base import BaseComponent
+from kotaemon.documents.base import Document, IO_Type
+from kotaemon.llms.chats.base import ChatLLM
+from kotaemon.llms.completions.base import LLM
+from kotaemon.prompt.base import BasePromptComponent
+
+
+class SimpleLinearPipeline(BaseComponent):
+    """
+    A simple pipeline for running a function with a prompt, a language model, and an
+        optional post-processor.
+
+    Attributes:
+        prompt (BasePromptComponent): The prompt component used to generate the initial
+            input.
+        llm (Union[ChatLLM, LLM]): The language model component used to generate the
+            output.
+        post_processor (Union[BaseComponent, Callable[[IO_Type], IO_Type]]): An optional
+            post-processor component or function.
+
+    Example Usage:
+        from kotaemon.llms.chats.openai import AzureChatOpenAI
+        from kotaemon.prompt.base import BasePromptComponent
+
+        def identity(x):
+            return x
+
+        llm = AzureChatOpenAI(
+            openai_api_base="your openai api base",
+            openai_api_key="your openai api key",
+            openai_api_version="your openai api version",
+            deployment_name="dummy-q2-gpt35",
+            temperature=0,
+            request_timeout=600,
+        )
+
+        pipeline = SimpleLinearPipeline(
+            prompt=BasePromptComponent(template="what is {word} in Japanese ?"),
+            llm=llm,
+            post_processor=identity,
+        )
+        print(pipeline(word="lone"))
+    """
+
+    prompt: BasePromptComponent
+    llm: Union[ChatLLM, LLM]
+    post_processor: Union[BaseComponent, Callable[[IO_Type], IO_Type]]
+
+    def run(
+        self,
+        *,
+        llm_kwargs: Optional[dict] = {},
+        post_processor_kwargs: Optional[dict] = {},
+        **prompt_kwargs,
+    ):
+        """
+        Run the function with the given arguments and return the final output as a
+            Document object.
+
+        Args:
+            llm_kwargs (dict): Keyword arguments for the llm call.
+            post_processor_kwargs (dict): Keyword arguments for the post_processor.
+            **prompt_kwargs: Keyword arguments for populating the prompt.
+
+        Returns:
+            Document: The final output of the function as a Document object.
+        """
+        prompt = self.prompt(**prompt_kwargs)
+        llm_output = self.llm(prompt.text, **llm_kwargs)
+        if self.post_processor is not None:
+            final_output = self.post_processor(llm_output, **post_processor_kwargs)
+        else:
+            final_output = llm_output
+
+        return Document(final_output)
+
+
+class GatedLinearPipeline(SimpleLinearPipeline):
+    """
+    A pipeline that extends the SimpleLinearPipeline class and adds a condition
+        attribute.
+
+    Attributes:
+        condition (Callable[[IO_Type], Any]): A callable function that represents the
+            condition.
+
+    Example Usage:
+        from kotaemon.llms.chats.openai import AzureChatOpenAI
+        from kotaemon.post_processing.extractor import RegexExtractor
+        from kotaemon.prompt.base import BasePromptComponent
+
+        def identity(x):
+            return x
+
+        llm = AzureChatOpenAI(
+            openai_api_base="your openai api base",
+            openai_api_key="your openai api key",
+            openai_api_version="your openai api version",
+            deployment_name="dummy-q2-gpt35",
+            temperature=0,
+            request_timeout=600,
+        )
+
+        pipeline = GatedLinearPipeline(
+            prompt=BasePromptComponent(template="what is {word} in Japanese ?"),
+            condition=RegexExtractor(pattern="some pattern"),
+            llm=llm,
+            post_processor=identity,
+        )
+        print(pipeline(condition_text="some pattern", word="lone"))
+        print(pipeline(condition_text="other pattern", word="lone"))
+    """
+
+    condition: Callable[[IO_Type], Any]
+
+    def run(
+        self,
+        *,
+        condition_text: Optional[str] = None,
+        llm_kwargs: Optional[dict] = {},
+        post_processor_kwargs: Optional[dict] = {},
+        **prompt_kwargs,
+    ) -> Document:
+        """
+        Run the pipeline with the given arguments and return the final output as a
+            Document object.
+
+        Args:
+            condition_text (str): The condition text to evaluate. Default to None.
+            llm_kwargs (dict): Additional keyword arguments for the language model call.
+            post_processor_kwargs (dict): Additional keyword arguments for the
+                post-processor.
+            **prompt_kwargs: Keyword arguments for populating the prompt.
+
+        Returns:
+            Document: The final output of the pipeline as a Document object.
+
+        Raises:
+            ValueError: If condition_text is None
+        """
+        if condition_text is None:
+            raise ValueError("`condition_text` must be provided")
+
+        if self.condition(condition_text):
+            return super().run(
+                llm_kwargs=llm_kwargs,
+                post_processor_kwargs=post_processor_kwargs,
+                **prompt_kwargs,
+            )
+
+        return Document(None)