Dynamic Tools

Author: paolorechia

.env.example

@@ -6,11 +6,11 @@ DOCUMENT_STORE_NAME=brainchulo_docs
 CONVERSATION_STORE_NAME=brainchulo_convos
 DATABASE_URL=sqlite:///data/brainchulo.db
 ANDROMEDA_URL=http://0.0.0.0:9000/
-
 MODEL_PATH=/models/vicuna-AlekseyKorshuk-7B-GPTQ-4bit-128g
-#  General Settings
+USE_FLOW_AGENTS=true
+# General Settings
 GENERAL_LOADING_METHOD=GPTQ
-GENERAL_BOOL_CPU_OFFLOADING=true
+GENERAL_BOOL_CPU_OFFLOADING=false
 # HuggingFace
 HF_BOOL_USE_QUANT=true
 HF_BOOL_USE_4_BIT=true

README.md

@@ -117,6 +117,12 @@ docker compose down
 docker-compose exec backend alembic upgrade head
 ```
 
+## Enabling Flow Agents
+We have a experimental feature that enables what we call flow agents. To activate it, set the feature flag to true in the .env file:
+```
+USE_FLOW_AGENTS=true
+```
+
 ## Creating a Plugin
 
 1. Create a new directory for your plugin, e.g. `/app/plugins/my_plugin`

app/agents/__init__.py

@@ -1,2 +1,2 @@
-from agents.document import DocumentQuestionAnswerAgent
-from agents.chain_of_thoughts import ChainOfThoughtsAgent
+from agents.chain_of_thoughts import ChainOfThoughtsAgent
+from agents.flow_cot import ChainOfThoughtsFlowAgent

app/agents/base.py

@@ -1,19 +1,20 @@
 # import the necessary libraries
 from abc import abstractmethod
 from colorama import Style, Fore
-from typing import Dict, Callable
+from typing import List
 from andromeda_chain import AndromedaChain
+from tools.base import BaseTool
 
 def color_print(msg, color):
     print(color + Style.BRIGHT + msg  + Style.RESET_ALL, flush=True)
 
 
 class BaseAgent:
     """Base Agent.
-    
+
     Nothing too exciting here.
     """
-    def __init__(self, andromeda: AndromedaChain, tools: Dict[str, Callable[[str], str]]):
+    def __init__(self, andromeda: AndromedaChain, tools: List[BaseTool]):
         self.andromeda = andromeda
         self.tools = tools
 
@@ -27,3 +28,22 @@ def do_tool(self, tool_name, act_input):
         color_print(f"Tool result: {result}", Fore.BLUE)
         return result
 
+    def _build_tool_description_line(self, tool: BaseTool):
+        return f"{tool.name}: {tool.short_description()}"
+
+    def prepare_start_prompt(self, prompt_start_template):
+        tools = [item for _, item in self.tools.items()]
+        tools_descriptions = "\n".join([self._build_tool_description_line(tool) for tool in tools])
+        action_list = str([tool.name for tool in tools]).replace('"', "")
+        few_shot_examples = "\n".join(
+            [
+                f"Example {idx}:\n{tool.few_shot_examples()}"
+                for idx, tool in enumerate(tools)
+            ]
+        )
+        self.valid_tools = [tool.name for tool in tools]
+        self.prepared_prompt = prompt_start_template.format(
+            tools_descriptions=tools_descriptions,
+            action_list=action_list,
+            few_shot_examples=few_shot_examples
+        )

app/agents/chain_of_thoughts.py

@@ -1,17 +1,10 @@
-from typing import Dict, Callable
-from andromeda_chain import AndromedaChain
-from agents.flow_based import BaseAgent, BaseFlowAgent
-from flow.flow import Flow, PromptNode, ToolNode, ChoiceNode, StartNode
-from prompts.flow_guidance_cot import FlowChainOfThoughts, PROMPT_START_STRING
+from agents.base import BaseAgent
 from guidance_tooling.guidance_programs.tools import ingest_file
 from guidance_tooling.guidance_programs.tools import clean_text
-from guidance_tooling.guidance_programs.tools import load_tools
 
 import os
 from langchain.llms import LlamaCpp
-from dotenv import load_dotenv
 import os
-from chromadb.config import Settings
 from colorama import Fore, Style
 from langchain.chains import RetrievalQA
 from langchain.llms import LlamaCpp
@@ -25,58 +18,15 @@ def get_llm():
     global llm
     if llm is None:
         print("Loading guidance model...")
-        model_type = "LlamaCpp"
         model_path ="/home/karajan/Downloads/airoboros-13b-gpt4.ggmlv3.q8_0.bin"
         model_n_ctx =1000
-        target_source_chunks = 4
         n_gpu_layers = 500
         use_mlock = 0
         n_batch = os.environ.get('N_BATCH') if os.environ.get('N_BATCH') != None else 512
         callbacks = []
-        qa_prompt = ""
         llm = LlamaCpp(model_path=model_path, n_ctx=model_n_ctx, callbacks=callbacks, verbose=False,n_gpu_layers=n_gpu_layers, use_mlock=use_mlock,top_p=0.9, n_batch=n_batch)
     return llm
-
-class ChainOfThoughtsFlowAgent(BaseFlowAgent):
-    def __init__(self, andromeda: AndromedaChain, tools: Dict[str, Callable[[str], str]]):
-        def execute_tool(variables):
-            action_name = variables["tool_name"]
-            action_input = variables["act_input"]
-            return self.do_tool(action_name, action_input)
-
-        start = StartNode("start", FlowChainOfThoughts.flow_prompt_start, {
-            "Action": "choose_action",
-            "Final Answer": "final_prompt"
-        })
-        thought = PromptNode("thought", FlowChainOfThoughts.thought_gen)
-        choose_action = PromptNode("choose_action", FlowChainOfThoughts.choose_action)
-        perform_action = PromptNode("perform_action", FlowChainOfThoughts.action_input)
-        execute_tool_node = ToolNode("execute_tool", execute_tool)
-        decide = ChoiceNode("decide", ["thought", "final_prompt"], max_decisions=3, force_exit_on="final_prompt")
-        final = PromptNode("final_prompt", FlowChainOfThoughts.final_prompt)
-
-        thought.set_next(choose_action)
-        choose_action.set_next(perform_action)
-        perform_action.set_next(execute_tool_node)
-        execute_tool_node.set_next(decide)
-
-        flow = Flow(
-            [start, thought, choose_action, perform_action, execute_tool_node, decide, final]
-        )
-
-        super().__init__(andromeda, tools, flow)
-        self.valid_tools = list(tools.keys())
-        self.valid_answers = ["Action", "Final Answer"]
-
-
-    def run(self, query: str) -> str:
-        return super().run(query, variables={
-            "prompt_start": PROMPT_START_STRING,
-            "question": query,
-            "valid_tools": self.valid_tools,
-            "valid_answers": self.valid_answers,
-        })
-
+ 
 class ChainOfThoughtsAgent(BaseAgent):
     def __init__(self, guidance, retriever, num_iter=3):
         self.guidance = guidance

app/agents/document.py

@@ -0,0 +1,49 @@
+from typing import Dict, Callable
+from andromeda_chain import AndromedaChain
+from agents.flow_based import BaseFlowAgent
+from flow.flow import Flow, PromptNode, ToolNode, ChoiceNode, StartNode
+from prompts.flow_guidance_cot import FlowChainOfThoughts, PROMPT_START_STRING
+
+
+class ChainOfThoughtsFlowAgent(BaseFlowAgent):
+    def __init__(self, andromeda: AndromedaChain, tools: Dict[str, Callable[[str], str]]):
+        def execute_tool(variables):
+            action_name = variables["tool_name"]
+            action_input = variables["act_input"]
+            return self.do_tool(action_name, action_input)
+
+        start = StartNode("start", FlowChainOfThoughts.flow_prompt_start, {
+            "Action": "choose_action",
+            "Final Answer": "final_prompt"
+        })
+        thought = PromptNode("thought", FlowChainOfThoughts.thought_gen)
+        choose_action = PromptNode("choose_action", FlowChainOfThoughts.choose_action)
+        perform_action = PromptNode("perform_action", FlowChainOfThoughts.action_input)
+        execute_tool_node = ToolNode("execute_tool", execute_tool)
+        decide = ChoiceNode("decide", ["thought", "final_prompt"], max_decisions=3, force_exit_on="final_prompt")
+        final = PromptNode("final_prompt", FlowChainOfThoughts.final_prompt)
+
+        thought.set_next(choose_action)
+        choose_action.set_next(perform_action)
+        perform_action.set_next(execute_tool_node)
+        execute_tool_node.set_next(decide)
+
+        flow = Flow(
+            [start, thought, choose_action, perform_action, execute_tool_node, decide, final]
+        )
+
+        super().__init__(andromeda, tools, flow)
+        self.valid_answers = ["Action", "Final Answer"]
+        self.prepare_start_prompt(PROMPT_START_STRING)
+
+
+    def run(self, query: str) -> str:
+        if not self.prepared_prompt:
+            raise TypeError("Your inherinted BaseFlowAgent class must call 'prepare_start_prompt' in it's constructor.")
+
+        return super().run(query, variables={
+            "prompt_start": self.prepared_prompt,
+            "question": query,
+            "valid_tools": self.valid_tools,
+            "valid_answers": self.valid_answers,
+        })

app/agents/flow_cot.py

@@ -3,7 +3,7 @@
 from langchain.text_splitter import RecursiveCharacterTextSplitter
 
 from andromeda_chain import AndromedaChain
-from agents import DocumentQuestionAnswerAgent, ChainOfThoughtsAgent
+from agents import ChainOfThoughtsAgent
 
 from settings import logger, load_config
 import guidance 
@@ -47,7 +47,6 @@ def __init__(self):
         }
         self.andromeda = AndromedaChain(config.andromeda_url)
         self.document_qa_agent = ChainOfThoughtsAgent(guidance, tools)
-        # self.document_qa_agent = DocumentQuestionAnswerAgent(self.andromeda, tools)
 
 
     def load_document(self, document_path, conversation_id=None):