modify prompts to improve entity and attribute value handling

src/agent/component_helpers.py

@@ -11,8 +11,8 @@
 from llama_index.core.output_parsers.utils import extract_json_str
 
 from prompt import (
-    ATTRIBUTE_IDENTIFICATION_PROMPT,
     ATTRIBUTE_CONDITION_PROMPT,
+    ATTRIBUTE_IDENTIFICATION_PROMPT,
     ENTITY_PROMPT,
     INDEX_PROMPT,
     INTENT_PROMPT,

src/agent/query_pipeline.py

@@ -3,8 +3,8 @@
 from agent.component_helpers import (
     construct_query,
     construct_url,
-    identify_attributes,
     define_attribute_condition,
+    identify_attributes,
     identify_entity,
     identify_index,
     identify_intent,

src/agent/test_query_pipeline.py

@@ -7,13 +7,13 @@
 from llama_index.llms.ollama import Ollama
 
 from agent.component_helpers import (
+    define_attribute_condition,
+    identify_attributes,
     identify_entity,
     identify_index,
     identify_intent,
     identify_rank,
-    identify_attributes,
     identify_time_frame,
-    define_attribute_condition,
 )
 from agent.goat_query_component import GoatQueryComponent
 
@@ -110,7 +110,8 @@
 
 
 @pytest.mark.parametrize(
-    "input_content, expected_index, _expected_intent, _expected_entities, _expected_rank, _expected_attribute, _expected_attribute_condition, _expected_time_from, _expected_time_to",
+    "input_content, expected_index, _expected_intent, _expected_entities, _expected_rank, _expected_attribute, "
+    "_expected_attribute_condition, _expected_time_from, _expected_time_to",
     QUERIES,
 )
 def test_index_module(
@@ -125,6 +126,7 @@ def test_index_module(
     _expected_time_to,
 ):
     # Skip the test if no expected index is provided
+    # sourcery skip: no-conditionals-in-tests
     if expected_index is None:
         pytest.skip("No expected index for this test case")
 
@@ -142,7 +144,8 @@ def test_index_module(
 
 
 @pytest.mark.parametrize(
-    "input_content, _expected_index, expected_intent, _expected_entities, _expected_rank, _expected_attribute, _expected_attribute_condition, _expected_time_from, _expected_time_to",
+    "input_content, _expected_index, expected_intent, _expected_entities, _expected_rank, _expected_attribute, "
+    "_expected_attribute_condition, _expected_time_from, _expected_time_to",
     QUERIES,
 )
 def test_intent_module(
@@ -157,6 +160,7 @@ def test_intent_module(
     _expected_time_to,
 ):
     # Skip the test if no expected intent is provided
+    # sourcery skip: no-conditionals-in-tests
     if expected_intent is None:
         pytest.skip("No expected intent for this test case")
 
@@ -174,7 +178,8 @@ def test_intent_module(
 
 
 @pytest.mark.parametrize(
-    "input_content, _expected_index, _expected_intent, expected_entities, _expected_rank, _expected_attribute, _expected_attribute_condition, _expected_time_from, _expected_time_to",
+    "input_content, _expected_index, _expected_intent, expected_entities, _expected_rank, _expected_attribute, "
+    "_expected_attribute_condition, _expected_time_from, _expected_time_to",
     QUERIES,
 )
 def test_entity_module(
@@ -189,6 +194,7 @@ def test_entity_module(
     _expected_time_to,
 ):
     # Skip the test if no expected entities are provided
+    # sourcery skip: no-conditionals-in-tests
     if expected_entities is None:
         pytest.skip("No expected entities for this test case")
 
@@ -204,6 +210,7 @@ def test_entity_module(
     # Assert the expected entity matches the result
     entities = result["state"]["entity"]["entities"]
     assert len(entities) == len(expected_entities)
+    # sourcery skip: no-loop-in-tests
     for entity in entities:
         assert (
             (entity["scientific_name"].lower() in [x.lower() for x in expected_entities])
@@ -213,7 +220,8 @@ def test_entity_module(
 
 
 @pytest.mark.parametrize(
-    "input_content, expected_index, expected_intent, expected_entities, expected_rank, _expected_attribute, _expected_attribute_condition, _expected_time_from, _expected_time_to",
+    "input_content, expected_index, expected_intent, expected_entities, expected_rank, _expected_attribute, "
+    "_expected_attribute_condition, _expected_time_from, _expected_time_to",
     QUERIES,
 )
 def test_rank_module(
@@ -228,6 +236,7 @@ def test_rank_module(
     _expected_time_to,
 ):
     # Skip the test if no expected rank is provided
+    # sourcery skip: no-conditionals-in-tests
     if expected_rank is None:
         pytest.skip("No expected rank for this test case")
 
@@ -259,7 +268,8 @@ def test_rank_module(
 
 
 @pytest.mark.parametrize(
-    "input_content, expected_index, _expected_intent, _expected_entities, _expected_rank, expected_attribute, _expected_attribute_condition, _expected_time_from, _expected_time_to",
+    "input_content, expected_index, _expected_intent, _expected_entities, _expected_rank, "
+    "expected_attribute, _expected_attribute_condition, _expected_time_from, _expected_time_to",
     QUERIES,
 )
 def test_attribute_module(
@@ -274,6 +284,7 @@ def test_attribute_module(
     _expected_time_to,
 ):
     # Skip the test if no expected attribute is provided
+    # sourcery skip: no-conditionals-in-tests
     if expected_attribute is None:
         pytest.skip("No expected attribute for this test case")
 
@@ -291,7 +302,8 @@ def test_attribute_module(
 
 
 @pytest.mark.parametrize(
-    "input_content, _expected_index, _expected_intent, _expected_entities, _expected_rank, _expected_attribute, _expected_attribute_condition, expected_time_from, expected_time_to",
+    "input_content, _expected_index, _expected_intent, _expected_entities, _expected_rank, "
+    "_expected_attribute, _expected_attribute_condition, expected_time_from, expected_time_to",
     QUERIES,
 )
 def test_time_module(
@@ -306,6 +318,7 @@ def test_time_module(
     expected_time_to,
 ):
     # Skip the test if no expected time is provided
+    # sourcery skip: no-conditionals-in-tests
     if expected_time_from is None and expected_time_to is None:
         pytest.skip("No expected time for this test case")
 
@@ -326,7 +339,8 @@ def test_time_module(
 
 
 @pytest.mark.parametrize(
-    "input_content, expected_index, _expected_intent, _expected_entities, _expected_rank, expected_attribute, expected_attribute_condition, _expected_time_from, _expected_time_to",
+    "input_content, expected_index, _expected_intent, _expected_entities, _expected_rank, "
+    "expected_attribute, expected_attribute_condition, _expected_time_from, _expected_time_to",
     QUERIES,
 )
 def test_define_attribute_condition(
@@ -341,6 +355,7 @@ def test_define_attribute_condition(
     _expected_time_to,
 ):
     # Skip the test if no expected attribute is provided
+    # sourcery skip: no-conditionals-in-tests
     if expected_attribute is None or expected_attribute_condition is None:
         pytest.skip("No expected attribute for this test case")
 

src/prompt.py

@@ -64,7 +64,8 @@
 
 **REMEMBER**: The query may refer to a subset of a larger group of organisms, for example in
 the query "What are the latest assemblies for the bivalve molluscs?", the taxon
-is "Bivalvia" and **NOT** "Mollusca" nor a lineage like "Mollusca: Bivalvia".
+is "Bivalvia" and **NOT** "Mollusca". If the taxon name contains multiple parts separated by ":",
+return only the part after the colon. E.g. for "Mollusca: Bivalvia", return only "Bivalvia".
 
 These queries have no taxon:
 - "what is the distribution of genome sizes across all classes?"
@@ -91,6 +92,11 @@
 If taxa is not applicable to the query, or the inferred scientific name is either None or an empty string,
 return an empty list.
 
+If entities are identified in the query, also generate a "masked_query" field in the response. This field
+should contain the query with each identified entity replaced with a placeholder. For example, if the query
+is "What are the latest assemblies for the families Canidae and Felidae?", the masked query would be "What are
+the latest assemblies for the families MASKED_ENTITY_0 and MASKED_ENTITY_1?".
+
 The query given by the user is as follows:
 `{query}`
 
@@ -103,10 +109,10 @@
             "scientific_name": "..."
         }}
     ],
+    "masked_query": "...",
     "explanation": "..."
 }}
 
-```json
 """
 )
 
@@ -155,7 +161,6 @@
     "explanation": "..."
 }}
 
-```json
 """
 )
 
@@ -188,7 +193,6 @@
     "explanation": "..."
 }}
 
-```json
 """
 )
 
@@ -232,7 +236,8 @@
 An intent can be one of the following three types:
 - **search**: The user is looking for a table of values.
 - **count**: The user is looking for a count of something.
-- **record**: The user is looking for a specific taxon, assembly or sample record without any condition or constraints.
+- **record**: The user is looking for a specific taxon, assembly, or sample record
+  without any condition or constraints.
 - **tree**: The user is looking for a phylogenetic tree.
 - **histogram**: The user is looking for a histogram chart showing the distribution of a **SINGLE** attribute.
 - **scatter**: The user is looking for a scatter plot chart showing the distribution of a pair of attributes.
@@ -265,7 +270,6 @@
     "explanation": "..."
 }}
 
-```json
 """
 )
 
@@ -314,55 +318,40 @@
 
 A user is trying to query a genomics database.
 
-We have identified some attributes in the query which might have specific conditions mentioned on them
-or they might simply be a "required" field in the output.
+We have identified some attributes in the query which may or may not have specific conditions mentioned on them.
+
+If the query can be interpreted as not placing any specific conditions on the value of the attribute, then the
+condition is "none".
 
 **Attribute List which needs to be filled with conditions:**
 `{attribute_metadata}`
 
 The query given by the user is as follows:
 `{query}`
 
-Your task is to return the above list of attributes and fill each attribute with conditions.
+Your task is to determine whether each attribute in the above list should have specific conditions attached or simply
+be required.
 The conditions can be one of the following:
->, <, >=, <=, =, !=, in, required
+>, <, >=, <=, =, !=, in, none
 
 If the condition is "in", the value will be a list of values.
-If the condition is "required", the value will be null, we just want the attribute to be "present" in the results.
+If the condition is "none", the value will be null, we just want the attribute to be **present** in the results.
 You need to reply in the following format:
 {{
     "attributes": [
         {{
             "attribute": "...",
-            "condition": "required",
+            "condition": "none",
             "value": "..." or ["...", "..."] or null
         }}
     ],
     "explanation": "why the condition was added..."
 }}
 
-e.g.
-query: "What is the contig N50 value for the family Canidae?"
-attributes: ["contig_n50"]
-
-response - 
-```json
-{
-    "attributes": [
-        {
-            "attribute": "contig_n50",
-            "condition": "required",
-            "value": null
-        }
-    ],
-    "explanation": "The contig N50 value is required in the output."
-}
-```
-
 *REMEMBER:*
-The result should contain ONLY attributes from the list given above, DO NOT add extra attributes.
-DO NOT omit any attributes from the list.
-In most cases, the condition will be "required".
+- The result should contain ONLY attributes from the list given above, DO NOT add extra attributes.
+- DO NOT omit any attributes from the list.
+- In case the exact condition is not specified, we can assume "none".
 
 """
 )
@@ -386,7 +375,8 @@
 `{query}`
 
 *IMPORTANT:*
-We do not want a programmatic answer, we only need the best matching taxon id which has the closest "rank" to the query.
+We do not want a programmatic answer, we only need the best matching taxon id
+which has the closest "rank" to the query.
 Do not give me python code, your response should simply be a JSON of the following format:
 {{
     "taxon_id": "...",
@@ -395,6 +385,5 @@
 
 The taxon_id HAS TO BE AN INTEGER.
 
-```json
 """
 )