[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/solver.py
+++ b/knowledgehub/pipelines/agents/rewoo/solver.py
@@ -0,0 +1,66 @@
+from typing import Any, List, Optional, Union
+
+from kotaemon.base import BaseComponent
+from kotaemon.prompt.template import PromptTemplate
+
+from ..base import BaseLLM
+from ..output.base import BaseScratchPad
+from .prompt import few_shot_solver_prompt, zero_shot_solver_prompt
+
+
+class Solver(BaseComponent):
+    model: BaseLLM
+    prompt_template: Optional[PromptTemplate] = None
+    examples: Optional[Union[str, List[str]]] = None
+
+    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, plan_evidence) -> str:
+        """
+        Compose the prompt from template, plan&evidence, examples and instruction.
+        """
+        fewshot = self._compose_fewshot_prompt()
+        if self.prompt_template is not None:
+            if "fewshot" in self.prompt_template.placeholders:
+                return self.prompt_template.populate(
+                    plan_evidence=plan_evidence, fewshot=fewshot, task=instruction
+                )
+            else:
+                return self.prompt_template.populate(
+                    plan_evidence=plan_evidence, task=instruction
+                )
+        else:
+            if self.examples is not None:
+                return few_shot_solver_prompt.populate(
+                    plan_evidence=plan_evidence, fewshot=fewshot, task=instruction
+                )
+            else:
+                return zero_shot_solver_prompt.populate(
+                    plan_evidence=plan_evidence, task=instruction
+                )
+
+    def run(
+        self,
+        instruction: str,
+        plan_evidence: str,
+        output: BaseScratchPad = BaseScratchPad(),
+    ) -> Any:
+        response = None
+        output.info("Running Solver")
+        output.debug(f"Instruction: {instruction}")
+        output.debug(f"Plan Evidence: {plan_evidence}")
+        prompt = self._compose_prompt(instruction, plan_evidence)
+        output.debug(f"Prompt: {prompt}")
+        try:
+            response = self.model(prompt)
+            output.info("Solver run successful.")
+        except ValueError:
+            output.error("Solver failed to retrieve response from LLM")
+
+        return response