Merge pull request #91 from gregdrizz/auto_agent_choosing

Auto agent choosing

agent/llm_utils.py

@@ -1,4 +1,7 @@
 from __future__ import annotations
+
+import json
+
 from fastapi import WebSocket
 import time
 
@@ -98,3 +101,77 @@ async def stream_response(model, messages, temperature, max_tokens, websocket):
                 paragraph = ""
     print(f"streaming response complete")
     return response
+
+
+def choose_agent(task: str) -> str:
+    """Determines what agent should be used
+    Args:
+        task (str): The research question the user asked
+    Returns:
+        agent - The agent that will be used
+        agent_role_prompt (str): The prompt for the agent
+    """
+    configuration = choose_agent_configuration()
+
+    response = openai.ChatCompletion.create(
+        model=CFG.smart_llm_model,
+        messages=[
+            {"role": "user", "content": f"{task}"}],
+        functions=configuration,
+        temperature=0,
+    )
+    message = response["choices"][0]["message"]
+
+    if message.get("function_call"):
+        function_name = message["function_call"]["name"]
+        return {"agent": json.loads(message["function_call"]["arguments"]).get("agent"),
+                "agent_role_prompt": json.loads(message["function_call"]["arguments"]).get("instructions")}
+    else:
+        {"agent": "Default Agent",
+         "agent_role_prompt": "You are an AI critical thinker research assistant. Your sole purpose is to write well written, critically acclaimed, objective and structured reports on given text."}
+
+
+def choose_agent_configuration():
+    configuration = [
+        {
+            "name": "research",
+            "description": "Researches the given topic even if it can't be answered",
+            "parameters": {
+                "type": "object",
+                "properties": {
+                    "agent": {
+                        "type": "string",
+                        "description":
+                            """
+                                Determines the field of the topic and the name of the agent we could use in order to research 
+                                about the topic provided.
+                                
+                                Example of agents:
+                                    "Business Analyst Agent", "Finance Agent", "Travel Agent",
+                                 "Academic Research Agent", "Computer Security Analyst Agent"
+                                 
+                                 if an agent for the field required doesn't exist make one up
+                            """,
+                    },
+                    "instructions": {
+                        "type": "string",
+                        "description":
+                            """
+                            each provided agent needs instructions in order to start working,
+                            examples for agents and their instructions:
+                                    "Finance Agent": "You are a seasoned finance analyst AI assistant. Your primary goal is to compose comprehensive, astute, impartial, and methodically arranged financial reports based on provided data and trends.",
+                                    "Travel Agent": "You are a world-travelled AI tour guide assistant. Your main purpose is to draft engaging, insightful, unbiased, and well-structured travel reports on given locations, including history, attractions, and cultural insights.",
+                                    "Academic Research Agent": "You are an AI academic research assistant. Your primary responsibility is to create thorough, academically rigorous, unbiased, and systematically organized reports on a given research topic, following the standards of scholarly work.",
+                                    "Business Analyst": "You are an experienced AI business analyst assistant. Your main objective is to produce comprehensive, insightful, impartial, and systematically structured business reports based on provided business data, market trends, and strategic analysis.",
+                                    "Computer Security Analyst Agent": "You are an AI specializing in computer security analysis. Your principal duty is to generate comprehensive, meticulously detailed, impartial, and systematically structured reports on computer security topics. This includes Exploits, Techniques, Threat Actors, and Advanced Persistent Threat (APT) Groups. All produced reports should adhere to the highest standards of scholarly work and provide in-depth insights into the complexities of computer security.",
+                                    
+                            """,
+                    },
+                },
+                "required": ["agent", "instructions"],
+            },
+        }
+    ]
+    return configuration
+
+

agent/research_agent.py

@@ -21,14 +21,15 @@
 
 
 class ResearchAgent:
-    def __init__(self, question, agent, websocket):
+    def __init__(self, question, agent, agent_role_prompt, websocket):
         """ Initializes the research assistant with the given question.
         Args: question (str): The question to research
         Returns: None
         """
 
         self.question = question
         self.agent = agent
+        self.agent_role_prompt = agent_role_prompt
         self.visited_urls = set()
         self.research_summary = ""
         self.directory_name = ''.join(c for c in question if c.isascii() and c not in string.punctuation)[:100]
@@ -68,7 +69,7 @@ async def get_new_urls(self, url_set_input):
     async def call_agent(self, action, stream=False, websocket=None):
         messages = [{
             "role": "system",
-            "content": prompts.generate_agent_role_prompt(self.agent),
+            "content": self.agent_role_prompt if self.agent_role_prompt else prompts.generate_agent_role_prompt(self.agent)
         }, {
             "role": "user",
             "content": action,

agent/run.py

@@ -34,19 +34,19 @@ async def disconnect(self, websocket: WebSocket):
         del self.sender_tasks[websocket]
         del self.message_queues[websocket]
 
-    async def start_streaming(self, task, report_type, agent, websocket):
-        report, path = await run_agent(task, report_type, agent, websocket)
+    async def start_streaming(self, task, report_type, agent, agent_role_prompt, websocket):
+        report, path = await run_agent(task, report_type, agent, agent_role_prompt, websocket)
         return report, path
 
 
-async def run_agent(task, report_type, agent, websocket):
+async def run_agent(task, report_type, agent, agent_role_prompt, websocket):
     check_openai_api_key()
 
     start_time = datetime.datetime.now()
 
     # await websocket.send_json({"type": "logs", "output": f"Start time: {str(start_time)}\n\n"})
 
-    assistant = ResearchAgent(task, agent, websocket)
+    assistant = ResearchAgent(task, agent, agent_role_prompt, websocket)
     await assistant.conduct_research()
 
     report, path = await assistant.write_report(report_type, websocket)