Query Embedding for Tool Rag

Author: EoghanOConnor

evaluator/algorithms/tool_rag_algorithm.py

@@ -549,6 +549,33 @@ async def process_query(self, query_spec: QuerySpecification) -> AlgoResponse:
         agent = create_react_agent(self._model, relevant_tools)
         return await self._invoke_agent_on_query(agent, query_spec.query), relevant_tool_names
 
+    def embed_additional_queries(self, query_specs: List[QuerySpecification]):
+        """
+        Embed all additional_queries found in the provided QuerySpecification instances.
+        Returns a dict: {query_id: {tool_name: [(queryN, embedding), ...]}}
+        """
+        if not hasattr(self, "embeddings") or self.embeddings is None:
+            raise RuntimeError("Embeddings must be initialized before calling this method.")
+        results = {}
+        for spec in query_specs:
+            if not getattr(spec, "additional_queries", None):
+                continue
+            spec_results = {}
+            add_queries = spec.additional_queries
+            # flatten if wrapped in {"additional_queries": ...}
+            if isinstance(add_queries, dict) and "additional_queries" in add_queries and len(add_queries) == 1:
+                add_queries = add_queries["additional_queries"]
+            for tool_name, queries in add_queries.items():
+                emb_list = []
+                for k, query in queries.items():
+                    if not isinstance(query, str):
+                        continue
+                    emb = self.embeddings.embed_query(query)
+                    emb_list.append((k, emb))
+                spec_results[tool_name] = emb_list
+            results[spec.id] = spec_results
+        return results
+
     def tear_down(self) -> None:
         connections.disconnect(alias=MILVUS_CONNECTION_ALIAS)
 

evaluator/components/data_provider.py

@@ -27,6 +27,8 @@ class QuerySpecification(BaseModel):
     """
     id: int
     query: str
+    additional_queries: Optional[Dict[str, Any]] = None
+    path: Optional[str] = None
     reference_answer: Optional[str] = None
     golden_tools: ToolSet = Field(default_factory=dict)
     additional_tools: Optional[ToolSet] = None
@@ -309,7 +311,7 @@ def _load_queries_from_single_file(
         root_dataset_path: str or Path,
         experiment_environment: EnvironmentConfig,
         dataset_config: DatasetConfig,
-) -> List[QuerySpecification]:
+) -> Tuple[List[QuerySpecification], List[Dict[str, Any]]]:
     with open(query_file_path, 'r') as f:
         data = json.load(f)
 
@@ -328,6 +330,13 @@ def _load_queries_from_single_file(
             log(f"Invalid query spec, skipping this query.")
         else:
             query = raw_query_spec.get("query")
+            if raw_query_spec.get("additional_queries"):
+                additional_queries = raw_query_spec.get("additional_queries")
+                print(f"Additional queries provided: {additional_queries}")
+
+            else:
+                print(f"No additional queries provided")
+                additional_queries = None
             query_id = int(raw_query_spec.get("query_id"))
             golden_tools, additional_tools = (
                 _parse_raw_query_tool_definitions(raw_query_spec, experiment_environment, dataset_config))
@@ -341,6 +350,8 @@ def _load_queries_from_single_file(
                     QuerySpecification(
                         id=query_id,
                         query=query,
+                        path=str(query_file_path),
+                        additional_queries=additional_queries,
                         reference_answer=reference_answer,
                         golden_tools=golden_tools,
                         additional_tools=additional_tools or None
@@ -358,7 +369,7 @@ def get_queries(
         experiment_environment: EnvironmentConfig,
         dataset_config: DatasetConfig,
         fine_tuning_mode=False
-) -> List[QuerySpecification]:
+) -> Tuple[List[QuerySpecification], List[Dict[str, Any]]]:
     """Load queries from the dataset."""
     root_dataset_path = Path(os.getenv("ROOT_DATASET_PATH"))
     if not root_dataset_path:
@@ -375,14 +386,14 @@ def get_queries(
     queries_num = None if fine_tuning_mode else dataset_config.queries_num
     queries = []
     for path in local_paths:
+        print(f"\n\n")
+        print(f"--------------------------------")
+        print(f"Loading queries from file: {path}")
+        print(f"\n\n")
         remaining_queries_num = None if queries_num is None else queries_num - len(queries)
         if remaining_queries_num == 0:
             break
-        new_queries = _load_queries_from_single_file(path,
-                                                     remaining_queries_num,
-                                                     root_dataset_path,
-                                                     experiment_environment,
-                                                     dataset_config)
+        new_queries= _load_queries_from_single_file(path, remaining_queries_num, root_dataset_path, experiment_environment, dataset_config)
         queries.extend(new_queries)
 
     return queries

evaluator/evaluator.py

@@ -3,7 +3,7 @@
 import time
 import traceback
 from typing import List, Tuple
-
+from pathlib import Path
 import openai
 from langgraph.errors import GraphRecursionError
 from pydantic import ValidationError
@@ -17,6 +17,8 @@
 from evaluator.interfaces.algorithm import Algorithm
 from evaluator.utils.csv_logger import CSVLogger
 from evaluator.components.llm_provider import get_llm
+from evaluator.utils.parsing_tools import generate_and_save_additional_queries
+import json as _json
 from dotenv import load_dotenv
 
 from evaluator.utils.tool_logger import ToolLogger
@@ -35,13 +37,13 @@ class Evaluator(object):
 
     config: EvaluationConfig
 
-    def __init__(self, config_path: str | None, use_defaults: bool):
+    def __init__(self, config_path: str | None, use_defaults: bool, test_with_additional_queries: bool = False):
         try:
             self.config = load_config(config_path, use_defaults=use_defaults)
         except ConfigError as ce:
             log(f"Configuration error: {ce}")
             raise SystemExit(2)
-
+        self.test_with_additional_queries = test_with_additional_queries
     async def run(self) -> None:
 
         # Set up the necessary components for the experiments:
@@ -112,15 +114,13 @@ async def _run_experiment(self,
         Runs the specified experiment and returns the number of evaluated queries.
         """
         processed_queries_num = 0
-
         try:
             queries = await self._set_up_experiment(spec, metric_collectors, mcp_proxy_manager)
             algorithm, environment = spec
 
             try:
                 for i, query_spec in enumerate(queries):
                     log(f"Processing query #{query_spec.id} (Experiment {exp_index} of {total_exp_num}, query {i+1} of {len(queries)})...")
-
                     for mc in metric_collectors:
                         mc.prepare_for_measurement(query_spec)
 
@@ -199,12 +199,12 @@ async def _set_up_experiment(self,
         log(f"Initializing LLM connection: {environment.model_id}")
         llm = get_llm(model_id=environment.model_id, model_config=self.config.models)
         log("Connection established successfully.\n")
-
         log("Fetching queries for the current experiment...")
         queries = get_queries(environment, self.config.data)
         log(f"Successfully loaded {len(queries)} queries.\n")
         print_iterable_verbose("The following queries will be executed:\n", queries)
-
+        log(f"Generating additional queries queries.\n")
+        generate_and_save_additional_queries(llm, queries)
         log("Retrieving tool definitions for the current experiment...")
         tool_specs = get_tools_from_queries(queries)
         tools = await mcp_proxy_manager.run_mcp_proxy(tool_specs, init_client=True).get_tools()
@@ -213,8 +213,40 @@ async def _set_up_experiment(self,
 
         log("Setting up the algorithm and the metric collectors...")
         algorithm.set_up(llm, tools)
+    
+        if algorithm.__module__ == "evaluator.algorithms.tool_rag_algorithm":
+            log("Embedding additional queries...")
+            additional_query_embeddings = algorithm.embed_additional_queries(queries)
+            print("Additional query embedding counts per query/tool:",
+                  {k: {tk: len(tv) for tk, tv in v.items()} for k, v in additional_query_embeddings.items()})
+        else:
+            print("No additional queries to embed.")
         for mc in metric_collectors:
             mc.set_up()
         log("All set!\n")
 
         return queries
+
+if __name__ == "__main__":
+    import argparse
+    parser = argparse.ArgumentParser(description="Run the Evaluator experiments.")
+    parser.add_argument("--config", type=str, default=None, help="Path to evaluation config YAML file")
+    parser.add_argument("--defaults", action="store_true", help="Use default config options if set")
+    parser.add_argument("--test-with-additional-queries", action="store_true", help="Test with additional queries")
+    args = parser.parse_args()
+
+    from evaluator.utils.utils import log
+
+    log("Starting Evaluator main...")
+    evaluator = Evaluator(
+        config_path=args.config,
+        use_defaults=args.defaults,
+        test_with_additional_queries=args.test_with_additional_queries
+    )
+    try:
+        import asyncio
+        asyncio.run(evaluator.run())
+        log("Evaluator finished successfully!")
+    except Exception as e:
+        log(f"Evaluator failed: {e}")
+        raise

evaluator/utils/parsing_tools.py

@@ -0,0 +1,48 @@
+from evaluator.components.llm_provider import query_llm
+from pathlib import Path
+import re
+import json
+from evaluator.utils.utils import print_iterable_verbose, log
+
+def generate_and_save_additional_queries(llm, queries):
+    """
+    For each query in queries, use the provided LLM to generate additional_queries if not present,
+    and save to the appropriate JSON file for that query (matching by query_id).
+    """
+
+    system_prompt = '''You create 5 additional queries for each tool and only return the additional queries information, given the query implemented, return in the following format as a JSON string:
+                {tool_name: {"query1": "", "query2": "", "query3": "", "query4": "", "query5": ""}}  '''
+    for i, query_spec in enumerate(queries):
+        # If additional_queries already present, skip generating and saving
+        if getattr(query_spec, 'additional_queries', None):
+            log(f"Skipping query_id {getattr(query_spec, 'id', '<N/A>')} because additional_queries is present.")
+            continue
+        path = Path(query_spec.path)
+        user_prompt = f"tool_name = {getattr(query_spec, 'golden_tools', {}).keys()}, Query= {getattr(query_spec, 'query', None)}"
+        result = query_llm(llm, system_prompt, user_prompt)
+        # Remove markdown/code block wrappers if present
+        match = re.search(r"</think>\s*(.*)", result, re.DOTALL)
+        result_text = match.group(1).strip() if match else result
+        # Try to extract the 'additional_queries' dict block
+        additional = None
+        result_text = result_text.strip()
+        try:
+            additional = json.loads(result_text)
+        except Exception as e:
+            additional = None
+        query_spec.additional_queries = additional
+        # Saving additional queries to the original query JSON file
+        if path and additional is not None:
+            if path.exists():
+                import json as _json
+                with path.open('r', encoding='utf-8') as f:
+                    orig_queries = _json.load(f)
+                for item in orig_queries:
+                    if (
+                        (item.get("query_id") == query_spec.id)
+                        or (str(item.get("query_id")) == str(query_spec.id))
+                    ):
+                        item["additional_queries"] = additional
+                with path.open('w', encoding='utf-8') as f:
+                    _json.dump(orig_queries, f, indent=2, ensure_ascii=False)
+                log(f"Successfully added additional queries to original file {path}")