Let query parameters take precedence for retriever tool (llamaindex/openai) (#183)

Signed-off-by: Ann Zhang <ann.zhang@databricks.com>

Author: annzhang-db

integrations/langchain/tests/unit_tests/test_vector_search_retriever_tool.py

@@ -330,3 +330,18 @@ def test_kwargs_are_passed_through() -> None:
         score_threshold=0.5,
         extra_param="something random",
     )
+
+
+def test_kwargs_override_both_num_results_and_query_type() -> None:
+    vector_search_tool = init_vector_search_tool(DELTA_SYNC_INDEX, num_results=10, query_type="ANN")
+    vector_search_tool._vector_store.similarity_search = MagicMock()
+
+    vector_search_tool.invoke(
+        {"query": "what cities are in Germany", "k": 3, "query_type": "HYBRID"},
+    )
+    vector_search_tool._vector_store.similarity_search.assert_called_once_with(
+        query="what cities are in Germany",
+        k=3,  # Should use overridden value
+        query_type="HYBRID",  # Should use overridden value
+        filter={},
+    )

integrations/llamaindex/src/databricks_llamaindex/vector_search_retriever_tool.py

@@ -117,15 +117,19 @@ def get_query_text_vector(query: str) -> Tuple[Optional[str], Optional[List[floa
             kwargs = {**kwargs, **(self.model_extra or {})}
             kwargs = {k: v for k, v in kwargs.items() if k in signature.parameters}
 
+            # Allow kwargs to override the default values upon invocation
+            num_results = kwargs.pop("num_results", self.num_results)
+            query_type = kwargs.pop("query_type", self.query_type)
+
             # Ensure that we don't have duplicate keys
             kwargs.update(
                 {
                     "query_text": query_text,
                     "query_vector": query_vector,
                     "columns": self.columns,
                     "filters": combined_filters,
-                    "num_results": self.num_results,
-                    "query_type": self.query_type,
+                    "num_results": num_results,
+                    "query_type": query_type,
                 }
             )
             search_resp = self._index.similarity_search(**kwargs)

integrations/llamaindex/tests/unit_tests/test_vector_search_retriever_tool.py

@@ -242,3 +242,18 @@ def test_filters_are_combined() -> None:
         query_type=vector_search_tool.query_type,
         query_vector=None,
     )
+
+
+def test_kwargs_override_both_num_results_and_query_type() -> None:
+    vector_search_tool = init_vector_search_tool(DELTA_SYNC_INDEX, num_results=10, query_type="ANN")
+    vector_search_tool._index = create_autospec(VectorSearchIndex, instance=True)
+
+    vector_search_tool.call(query="what cities are in Germany", num_results=3, query_type="HYBRID")
+    vector_search_tool._index.similarity_search.assert_called_once_with(
+        columns=vector_search_tool.columns,
+        query_text="what cities are in Germany",
+        filters={},
+        num_results=3,  # Should use overridden value
+        query_type="HYBRID",  # Should use overridden value
+        query_vector=None,
+    )

integrations/openai/src/databricks_openai/vector_search_retriever_tool.py

@@ -232,14 +232,19 @@ def execute(
         signature = inspect.signature(self._index.similarity_search)
         kwargs = {**kwargs, **(self.model_extra or {})}
         kwargs = {k: v for k, v in kwargs.items() if k in signature.parameters}
+
+        # Allow kwargs to override the default values upon invocation
+        num_results = kwargs.pop("num_results", self.num_results)
+        query_type = kwargs.pop("query_type", self.query_type)
+
         kwargs.update(
             {
                 "query_text": query_text,
                 "query_vector": query_vector,
                 "columns": self.columns,
                 "filters": combined_filters,
-                "num_results": self.num_results,
-                "query_type": self.query_type,
+                "num_results": num_results,
+                "query_type": query_type,
             }
         )
         search_resp = self._index.similarity_search(**kwargs)

integrations/openai/tests/unit_tests/test_vector_search_retriever_tool.py

@@ -363,3 +363,20 @@ def test_filters_are_combined() -> None:
         query_type=vector_search_tool.query_type,
         query_vector=None,
     )
+
+
+def test_kwargs_override_both_num_results_and_query_type() -> None:
+    vector_search_tool = init_vector_search_tool(DELTA_SYNC_INDEX, num_results=10, query_type="ANN")
+    vector_search_tool._index = create_autospec(VectorSearchIndex, instance=True)
+
+    vector_search_tool.execute(
+        query="what cities are in Germany", num_results=3, query_type="HYBRID"
+    )
+    vector_search_tool._index.similarity_search.assert_called_once_with(
+        columns=vector_search_tool.columns,
+        query_text="what cities are in Germany",
+        filters={},
+        num_results=3,  # Should use overridden value
+        query_type="HYBRID",  # Should use overridden value
+        query_vector=None,
+    )