Fix resolving of references in ChatVertexAI.with_structured_output (#843)

Author: SebastianThomas1

libs/vertexai/langchain_google_vertexai/chat_models.py

@@ -1150,12 +1150,12 @@ class Joke(BaseModel):
 
     logprobs: Union[bool, int] = False
     """Whether to return logprobs as part of AIMessage.response_metadata.
-    
-    If False, don't return logprobs. If True, return logprobs for top candidate. 
+
+    If False, don't return logprobs. If True, return logprobs for top candidate.
     If int, return logprobs for top ``logprobs`` candidates.
-    
+
     **NOTE**: As of 10.28.24 this is only supported for gemini-1.5-flash models.
-    
+
     .. versionadded: 2.0.6
     """
     labels: Optional[Dict[str, str]] = None
@@ -2025,41 +2025,43 @@ class AnswerWithJustification(BaseModel):
         parser: OutputParserLike
 
         if method == "json_mode":
-            schema_is_typeddict = is_typeddict(schema)
-            if isinstance(schema, type) and not schema_is_typeddict:
-                # TODO: This gets the json schema of a pydantic model. It fails for
-                # nested models because the generated schema contains $refs that the
-                # gemini api doesn't support. We can implement a postprocessing function
-                # that takes care of this if necessary.
+            if isinstance(schema, type) and is_basemodel_subclass(schema):
                 if issubclass(schema, BaseModelV1):
                     schema_json = schema.schema()
                 else:
                     schema_json = schema.model_json_schema()
-                schema_json = replace_defs_in_schema(schema_json)
                 parser = PydanticOutputParser(pydantic_object=schema)
             else:
-                if schema_is_typeddict:
+                if is_typeddict(schema):
                     schema_json = convert_to_json_schema(schema)
+                elif isinstance(schema, dict):
+                    schema_json = schema
                 else:
-                    schema_json = cast(dict, schema)
+                    raise ValueError(f"Unsupported schema type {type(schema)}")
                 parser = JsonOutputParser()
+
+            # Resolve refs in schema because they are not supported
+            # by the Gemini API.
+            schema_json = replace_defs_in_schema(schema_json)
+
             llm = self.bind(
                 response_mime_type="application/json",
                 response_schema=schema_json,
                 ls_structured_output_format={
                     "kwargs": {"method": method},
-                    "schema": convert_to_json_schema(schema),
+                    "schema": schema_json,
                 },
             )
-
         else:
             tool_name = _get_tool_name(schema)
             if isinstance(schema, type) and is_basemodel_subclass(schema):
                 parser = PydanticToolsParser(tools=[schema], first_tool_only=True)
-            else:
+            elif is_typeddict(schema) or isinstance(schema, dict):
                 parser = JsonOutputKeyToolsParser(
                     key_name=tool_name, first_tool_only=True
                 )
+            else:
+                raise ValueError(f"Unsupported schema type {type(schema)}")
             tool_choice = tool_name if self._is_gemini_advanced else None
 
             try:

libs/vertexai/langchain_google_vertexai/functions_utils.py

@@ -112,7 +112,7 @@ def _format_json_schema_to_gapic(
     """Format a JSON schema from a Pydantic V2 BaseModel to gapic."""
     converted_schema: Dict[str, Any] = {}
     for key, value in schema.items():
-        if key == "definitions":
+        if key == "$defs":
             continue
         elif key == "items":
             converted_schema["items"] = _format_json_schema_to_gapic(
@@ -157,7 +157,10 @@ def _format_json_schema_to_gapic(
 def _dict_to_gapic_schema(
     schema: Dict[str, Any], pydantic_version: str = "v1"
 ) -> gapic.Schema:
+    # Resolve refs in schema because $refs and $defs are not supported
+    # by the Gemini API.
     dereferenced_schema = dereference_refs(schema)
+
     if pydantic_version == "v1":
         formatted_schema = _format_json_schema_to_gapic_v1(dereferenced_schema)
     else: