chore: Support the query method to handle types.ChatResponse response in Llama index template.

PiperOrigin-RevId: 741312199

Author: shawn-yang-google

tests/unit/vertex_llama_index/test_reasoning_engine_templates_llama_index.py

@@ -13,19 +13,20 @@
 # limitations under the License.
 #
 import importlib
-from unittest import mock
 import json
+from unittest import mock
 
 from google import auth
 import vertexai
 from google.cloud.aiplatform import initializer
 from vertexai.preview.reasoning_engines.templates import llama_index
 from vertexai.reasoning_engines import _utils
-import pytest
 
 from llama_index.core import prompts
 from llama_index.core.base.llms import types
 
+import pytest
+
 _TEST_LOCATION = "us-central1"
 _TEST_PROJECT = "test-project"
 _TEST_MODEL = "gemini-1.0-pro"
@@ -232,3 +233,88 @@ def test_enable_tracing_warning(self, caplog, llama_index_instrumentor_none_mock
  # TODO(b/384730642): Re-enable this test once the parent issue is fixed.
  # agent.set_up()
  # assert "enable_tracing=True but proceeding with tracing disabled" in caplog.text
+
+
+class TestToJsonSerializableLlamaIndexObject:
+ """Tests for `_utils.to_json_serializable_llama_index_object`."""
+
+ def test_llama_index_response(self):
+ mock_response: _utils.LlamaIndexResponse = mock.Mock(
+ spec=_utils.LlamaIndexResponse
+ )
+ mock_response.response = "test response"
+ mock_response.source_nodes = [
+ mock.Mock(
+ spec=_utils.LlamaIndexBaseModel,
+ model_dump_json=lambda: '{"name": "model1"}',
+ ),
+ mock.Mock(
+ spec=_utils.LlamaIndexBaseModel,
+ model_dump_json=lambda: '{"name": "model2"}',
+ ),
+ ]
+ mock_response.metadata = {"key": "value"}
+
+ want = {
+ "response": "test response",
+ "source_nodes": ['{"name": "model1"}', '{"name": "model2"}'],
+ "metadata": {"key": "value"},
+ }
+ got = _utils.to_json_serializable_llama_index_object(mock_response)
+ assert got == want
+
+ def test_llama_index_chat_response(self):
+ mock_chat_response: _utils.LlamaIndexChatResponse = mock.Mock(
+ spec=_utils.LlamaIndexChatResponse
+ )
+ mock_chat_response.message = mock.Mock(
+ spec=_utils.LlamaIndexBaseModel,
+ model_dump_json=lambda: '{"content": "chat message"}',
+ )
+
+ want = {"content": "chat message"}
+ got = _utils.to_json_serializable_llama_index_object(mock_chat_response)
+ assert got == want
+
+ def test_llama_index_base_model(self):
+ mock_base_model: _utils.LlamaIndexBaseModel = mock.Mock(
+ spec=_utils.LlamaIndexBaseModel
+ )
+ mock_base_model.model_dump_json = lambda: '{"name": "test_model"}'
+
+ want = {"name": "test_model"}
+ got = _utils.to_json_serializable_llama_index_object(mock_base_model)
+ assert got == want
+
+ def test_sequence_of_llama_index_base_model(self):
+ mock_base_model1: _utils.LlamaIndexBaseModel = mock.Mock(
+ spec=_utils.LlamaIndexBaseModel
+ )
+ mock_base_model1.model_dump_json = lambda: '{"name": "test_model1"}'
+ mock_base_model2: _utils.LlamaIndexBaseModel = mock.Mock(
+ spec=_utils.LlamaIndexBaseModel
+ )
+ mock_base_model2.model_dump_json = lambda: '{"name": "test_model2"}'
+ mock_base_model_list = [mock_base_model1, mock_base_model2]
+
+ want = [{"name": "test_model1"}, {"name": "test_model2"}]
+ got = _utils.to_json_serializable_llama_index_object(mock_base_model_list)
+ assert got == want
+
+ def test_sequence_of_mixed_types(self):
+ mock_base_model: _utils.LlamaIndexBaseModel = mock.Mock(
+ spec=_utils.LlamaIndexBaseModel
+ )
+ mock_base_model.model_dump_json = lambda: '{"name": "test_model"}'
+ mock_string = "test_string"
+ mock_list = [mock_base_model, mock_string]
+
+ want = [{"name": "test_model"}, "test_string"]
+ got = _utils.to_json_serializable_llama_index_object(mock_list)
+ assert got == want
+
+ def test_other_type(self):
+ test_dict = {"name": "test_model"}
+ want = "{'name': 'test_model'}"
+ got = _utils.to_json_serializable_llama_index_object(test_dict)
+ assert got == want

vertexai/preview/reasoning_engines/templates/llama_index.py

@@ -523,7 +523,7 @@ def query(
  self,
  input: Union[str, Mapping[str, Any]],
  **kwargs: Any,
- ) -> Union[Dict[str, Any], Sequence[Dict[str, Any]]]:
+ ) -> Union[str, Dict[str, Any], Sequence[Union[str, Dict[str, Any]]]]:
  """Queries the Agent with the given input and config.
 
  Args:
@@ -536,19 +536,14 @@ def query(
  Returns:
  The output of querying the Agent with the given input and config.
  """
- import json
  from vertexai.reasoning_engines import _utils
- from llama_index.core.base.response import schema
 
  if isinstance(input, str):
  input = {"input": input}
 
  if not self._runnable:
  self.set_up()
 
- response = self._runnable.run(**input, **kwargs)
- if isinstance(response, schema.Response):
- return _utils.llama_index_response_to_dict(response)
- if isinstance(response, Sequence):
- return [json.loads(r.model_dump_json()) for r in response]
- return json.loads(response.model_dump_json())
+ return _utils.to_json_serializable_llama_index_object(
+ self._runnable.run(**input, **kwargs)
+ )

vertexai/reasoning_engines/_utils.py

@@ -39,10 +39,15 @@
 
 try:
  from llama_index.core.base.response import schema as llama_index_schema
+ from llama_index.core.base.llms import types as llama_index_types
 
  LlamaIndexResponse = llama_index_schema.Response
+ LlamaIndexBaseModel = llama_index_schema.BaseModel
+ LlamaIndexChatResponse = llama_index_types.ChatResponse
 except ImportError:
  LlamaIndexResponse = Any
+ LlamaIndexBaseModel = Any
+ LlamaIndexChatResponse = Any
 
 JsonDict = Dict[str, Any]
 
@@ -111,7 +116,7 @@ def dataclass_to_dict(obj: dataclasses.dataclass) -> JsonDict:
  return json.loads(json.dumps(dataclasses.asdict(obj)))
 
 
-def llama_index_response_to_dict(obj: LlamaIndexResponse) -> Dict[str, Any]:
+def _llama_index_response_to_dict(obj: LlamaIndexResponse) -> Dict[str, Any]:
  response = {}
  if hasattr(obj, "response"):
  response["response"] = obj.response
@@ -123,6 +128,44 @@ def llama_index_response_to_dict(obj: LlamaIndexResponse) -> Dict[str, Any]:
  return json.loads(json.dumps(response))
 
 
+def _llama_index_chat_response_to_dict(
+ obj: LlamaIndexChatResponse,
+) -> Dict[str, Any]:
+ return json.loads(obj.message.model_dump_json())
+
+
+def _llama_index_base_model_to_dict(
+ obj: LlamaIndexBaseModel,
+) -> Dict[str, Any]:
+ return json.loads(obj.model_dump_json())
+
+
+def to_json_serializable_llama_index_object(
+ obj: Union[
+ LlamaIndexResponse,
+ LlamaIndexBaseModel,
+ LlamaIndexChatResponse,
+ Sequence[LlamaIndexBaseModel],
+ ]
+) -> Union[str, Dict[str, Any], Sequence[Union[str, Dict[str, Any]]]]:
+ """Converts a LlamaIndexResponse to a JSON serializable object."""
+ if isinstance(obj, LlamaIndexResponse):
+ return _llama_index_response_to_dict(obj)
+ if isinstance(obj, LlamaIndexChatResponse):
+ return _llama_index_chat_response_to_dict(obj)
+ if isinstance(obj, Sequence):
+ seq_result = []
+ for item in obj:
+ if isinstance(item, LlamaIndexBaseModel):
+ seq_result.append(_llama_index_base_model_to_dict(item))
+ continue
+ seq_result.append(str(item))
+ return seq_result
+ if isinstance(obj, LlamaIndexBaseModel):
+ return _llama_index_base_model_to_dict(obj)
+ return str(obj)
+
+
 def yield_parsed_json(body: httpbody_pb2.HttpBody) -> Iterable[Any]:
  """Converts the contents of the httpbody message to JSON format.