FS-76 Update agent selection prompt (#45)

.env.example

@@ -41,11 +41,9 @@ INTENT_AGENT_LLM="openai"
 REPORT_AGENT_LLM="mistral"
 VALIDATOR_AGENT_LLM="openai"
 DATASTORE_AGENT_LLM="openai"
-MATHS_AGENT_LLM="openai"
 WEB_AGENT_LLM="openai"
 CHART_GENERATOR_LLM="openai"
 ROUTER_LLM="openai"
-FILE_AGENT_LLM="openai"
 SUGGESTIONS_LLM="openai"
 DYNAMIC_KNOWLEDGE_GRAPH_LLM="openai"
 
@@ -56,10 +54,8 @@ INTENT_AGENT_MODEL="gpt-4o-mini"
 REPORT_AGENT_MODEL="mistral-large-latest"
 VALIDATOR_AGENT_MODEL="gpt-4o-mini"
 DATASTORE_AGENT_MODEL="gpt-4o-mini"
-MATHS_AGENT_MODEL="gpt-4o-mini"
 WEB_AGENT_MODEL="gpt-4o-mini"
 CHART_GENERATOR_MODEL="gpt-4o-mini"
 ROUTER_MODEL="gpt-4o-mini"
-FILE_AGENT_MODEL="gpt-4o-mini"
 SUGGESTIONS_MODEL="gpt-4o-mini"
 DYNAMIC_KNOWLEDGE_GRAPH_MODEL="gpt-4o"

backend/promptfoo/agent_selection_config.yaml

@@ -0,0 +1,85 @@
+description: "Agent selection"
+
+define:
+  &list_of_agents [
+    {
+      "name": "DatastoreAgent",
+      "description": "This agent is responsible for handling database queries to the bloomberg.csv dataset. This includes retrieving ESG scores, financial metrics, and other bloomberg-specific information. It interacts with the graph database to extract, process, and return ESG-related information from various sources, such as company sustainability reports or fund portfolios. This agent can not complete any task that is not specifically about the bloomberg.csv dataset.",
+    },
+    {
+      "name": "WebAgent",
+      "description": "This agent can perform general internet searches to complete the task by retrieving and summarizing the results and it can also perform web scrapes to retreive specific inpormation from web pages.",
+    },
+    {
+      "name": "ChartGeneratorAgent",
+      "description": "This agent is responsible for creating charts",
+    },
+  ]
+
+providers:
+  - id: openai:gpt-4o-mini
+    config:
+      temperature: 0
+
+prompts: file://promptfoo_test_runner.py:create_prompt
+
+tests:
+  - description: "Test the WebAgent is selected for an ESG data query when not related to bloomberg dataset"
+    vars:
+      user_prompt_template: "agent-selection-user-prompt"
+      user_prompt_args:
+        task: "What are the environmental factors that influence a company's ESG score?"
+      system_prompt_template: "agent-selection-system-prompt"
+      system_prompt_args:
+        list_of_agents: *list_of_agents
+    assert:
+      - type: javascript
+        value: JSON.parse(output).agent_name === "WebAgent"
+
+  - description: "Test the WebAgent is selected for queries about the news"
+    vars:
+      user_prompt_template: "agent-selection-user-prompt"
+      user_prompt_args:
+        task: "Are there any recent news articles discussing 3M Co's ESG initiatives?"
+      system_prompt_template: "agent-selection-system-prompt"
+      system_prompt_args:
+        list_of_agents: *list_of_agents
+    assert:
+      - type: javascript
+        value: JSON.parse(output).agent_name === "WebAgent"
+
+  - description: "Test the DatastoreAgent is selected for a data query related to the bloomberg dataset"
+    vars:
+      user_prompt_template: "agent-selection-user-prompt"
+      user_prompt_args:
+        task: "Which company name has the highest environmental ESG score in the bloomberg dataset?"
+      system_prompt_template: "agent-selection-system-prompt"
+      system_prompt_args:
+        list_of_agents: *list_of_agents
+    assert:
+      - type: javascript
+        value: JSON.parse(output).agent_name === "DatastoreAgent"
+
+  - description: "Test the DatastoreAgent is selected for a data query from the bloomberg.csv"
+    vars:
+      user_prompt_template: "agent-selection-user-prompt"
+      user_prompt_args:
+        task: "From the Bloomberg.csv data show me all the governance scores for american airlines in date order starting in the past."
+      system_prompt_template: "agent-selection-system-prompt"
+      system_prompt_args:
+        list_of_agents: *list_of_agents
+    assert:
+      - type: javascript
+        value: JSON.parse(output).agent_name === "DatastoreAgent"
+
+  - description: "Test the DatastoreAgent is selected for a query of 'the database'"
+    vars:
+      user_prompt_template: "agent-selection-user-prompt"
+      user_prompt_args:
+        task: "Check the database and give me all data for Masco Corp"
+      system_prompt_template: "agent-selection-system-prompt"
+      system_prompt_args:
+        list_of_agents: *list_of_agents
+    assert:
+      - type: javascript
+        value: JSON.parse(output).agent_name === "DatastoreAgent"

backend/src/agents/__init__.py

@@ -9,7 +9,6 @@
 from src.agents.validator_agent import ValidatorAgent
 from src.agents.answer_agent import AnswerAgent
 from src.agents.chart_generator_agent import ChartGeneratorAgent
-from src.agents.file_agent import FileAgent
 from src.agents.report_agent import ReportAgent
 
 
@@ -37,14 +36,11 @@ def agent_details(agent) -> dict:
 
 
 def get_available_agents() -> List[Agent]:
-    return [DatastoreAgent(config.datastore_agent_llm, config.datastore_agent_model),
-            WebAgent(config.web_agent_llm, config.web_agent_model),
-            ChartGeneratorAgent(config.chart_generator_llm,
-                                config.chart_generator_model),
-            FileAgent(config.file_agent_llm, config.file_agent_model),
-            # FS-63 Silencing Math agent - tool is not optimised.
-            # MathsAgent(config.maths_agent_llm, config.maths_agent_model),
-            ]
+    return [
+        DatastoreAgent(config.datastore_agent_llm, config.datastore_agent_model),
+        WebAgent(config.web_agent_llm, config.web_agent_model),
+        ChartGeneratorAgent(config.chart_generator_llm, config.chart_generator_model),
+    ]
 
 
 def get_agent_details():

backend/src/agents/datastore_agent.py

@@ -17,10 +17,10 @@
 
 cache = {}
 
+
 async def generate_cypher_query_core(
     question_intent, operation, question_params, aggregation, sort_order, timeframe, llm: LLM, model
 ) -> str:
-
     details_to_create_cypher_query = engine.load_prompt(
         "details-to-create-cypher-query",
         question_intent=question_intent,
@@ -38,8 +38,9 @@ async def generate_cypher_query_core(
             "generate-cypher-query", graph_schema=graph_schema, current_date=datetime.now()
         )
 
-        llm_query = await llm.chat(model, generate_cypher_query_prompt, details_to_create_cypher_query,
-                                   return_json=True)
+        llm_query = await llm.chat(
+            model, generate_cypher_query_prompt, details_to_create_cypher_query, return_json=True
+        )
         json_query = to_json(llm_query)
         await publish_log_info(LogPrefix.USER, f"Cypher generated by the LLM: {llm_query}", __name__)
         if json_query["query"] == "None":
@@ -49,12 +50,10 @@ async def generate_cypher_query_core(
     except Exception as e:
         logger.error(f"Error during data retrieval: {e}")
         raise
-    response = {
-        "content": db_response,
-        "ignore_validation": "false"
-    }
+    response = {"content": db_response, "ignore_validation": "false"}
     return json.dumps(response, indent=4)
 
+
 @tool(
     name="generate cypher query",
     description="Generate Cypher query if the category is data driven, based on the operation to be performed",
@@ -88,11 +87,12 @@ async def generate_cypher_query_core(
         ),
     },
 )
-
-async def generate_cypher(question_intent, operation, question_params, aggregation, sort_order,
-                          timeframe, llm: LLM, model) -> str:
-    return await generate_cypher_query_core(question_intent, operation, question_params, aggregation, sort_order,
-                                            timeframe, llm, model)
+async def generate_cypher(
+    question_intent, operation, question_params, aggregation, sort_order, timeframe, llm: LLM, model
+) -> str:
+    return await generate_cypher_query_core(
+        question_intent, operation, question_params, aggregation, sort_order, timeframe, llm, model
+    )
 
 
 async def get_semantic_layer_cache(llm, model, graph_schema):
@@ -104,14 +104,10 @@ async def get_semantic_layer_cache(llm, model, graph_schema):
     else:
         return cache
 
+
 @agent(
     name="DatastoreAgent",
-    description=(
-        "This agent is responsible for handling database queries related to ESG (Environmental, Social, Governance) "
-        "data, including retrieving ESG scores, financial metrics, and other company-specific or fund-specific "
-        "information. It interacts with the graph database to extract, process, and return ESG-related information "
-        "from various sources, such as company sustainability reports or fund portfolios."
-    ),
+    description="This agent is responsible for handling database queries to the bloomberg.csv dataset. This includes retrieving ESG scores, financial metrics, and other bloomberg-specific information. It interacts with the graph database to extract, process, and return ESG-related information from various sources, such as company sustainability reports or fund portfolios. This agent can not complete any task that is not specifically about the bloomberg.csv dataset.",  # noqa: E501
     tools=[generate_cypher],
 )
 class DatastoreAgent(Agent):