[AUR-431, AUR-435] Add Agent Interface and ReWOO Agent implementation (#31)

* add base Tool * minor update test_tool * update test dependency * update test dependency * Fix namespace conflict * update test * add base Agent Interface, add ReWoo Agent * minor update * update test * fix typo * remove unneeded print * update rewoo agent --------- Co-authored-by: trducng <trungduc1992@gmail.com>
2023-10-01 11:53:08 +07:00
parent f9fc02a32a
commit 91048770fa
11 changed files with 919 additions and 0 deletions
--- a/knowledgehub/pipelines/agents/rewoo/planner.py
+++ b/knowledgehub/pipelines/agents/rewoo/planner.py
@@ -0,0 +1,82 @@
+from typing import Any, List, Optional, Union
+
+from kotaemon.base import BaseComponent
+from kotaemon.prompt.template import PromptTemplate
+
+from ..base import BaseLLM, BaseTool
+from ..output.base import BaseScratchPad
+from .prompt import zero_shot_planner_prompt
+
+
+class Planner(BaseComponent):
+    model: BaseLLM
+    prompt_template: Optional[PromptTemplate] = None
+    examples: Optional[Union[str, List[str]]] = None
+    plugins: List[BaseTool]
+
+    def _compose_worker_description(self) -> str:
+        """
+        Compose the worker prompt from the workers.
+
+        Example:
+        toolname1[input]: tool1 description
+        toolname2[input]: tool2 description
+        """
+        prompt = ""
+        try:
+            for worker in self.plugins:
+                prompt += f"{worker.name}[input]: {worker.description}\n"
+        except Exception:
+            raise ValueError("Worker must have a name and description.")
+        return prompt
+
+    def _compose_fewshot_prompt(self) -> str:
+        if self.examples is None:
+            return ""
+        if isinstance(self.examples, str):
+            return self.examples
+        else:
+            return "\n\n".join([e.strip("\n") for e in self.examples])
+
+    def _compose_prompt(self, instruction) -> str:
+        """
+        Compose the prompt from template, worker description, examples and instruction.
+        """
+        worker_desctription = self._compose_worker_description()
+        fewshot = self._compose_fewshot_prompt()
+        if self.prompt_template is not None:
+            if "fewshot" in self.prompt_template.placeholders:
+                return self.prompt_template.populate(
+                    tool_description=worker_desctription,
+                    fewshot=fewshot,
+                    task=instruction,
+                )
+            else:
+                return self.prompt_template.populate(
+                    tool_description=worker_desctription, task=instruction
+                )
+        else:
+            if self.examples is not None:
+                return zero_shot_planner_prompt.populate(
+                    tool_description=worker_desctription,
+                    fewshot=fewshot,
+                    task=instruction,
+                )
+            else:
+                return zero_shot_planner_prompt.populate(
+                    tool_description=worker_desctription, task=instruction
+                )
+
+    def run(self, instruction: str, output: BaseScratchPad = BaseScratchPad()) -> Any:
+        response = None
+        output.info("Running Planner")
+        prompt = self._compose_prompt(instruction)
+        output.debug(f"Prompt: {prompt}")
+        try:
+            response = self.model(prompt)
+            output.info("Planner run successful.")
+        except ValueError:
+            output.error("Planner failed to retrieve response from LLM")
+            raise ValueError("Planner failed to retrieve response from LLM")
+
+        return response