From e04227668dc86a45183fa39357dce55318c161cf Mon Sep 17 00:00:00 2001 From: yxia216 Date: Sat, 24 May 2025 00:11:04 -0400 Subject: [PATCH 1/2] support-json-mode --- .../langchain_google_vertexai/chat_models.py | 24 +++++++++++++++---- .../functions_utils.py | 11 +++++---- 2 files changed, 25 insertions(+), 10 deletions(-) diff --git a/libs/vertexai/langchain_google_vertexai/chat_models.py b/libs/vertexai/langchain_google_vertexai/chat_models.py index 391e53d8e..c1aefc45f 100644 --- a/libs/vertexai/langchain_google_vertexai/chat_models.py +++ b/libs/vertexai/langchain_google_vertexai/chat_models.py @@ -143,7 +143,10 @@ from pydantic.v1 import BaseModel as BaseModelV1 from typing_extensions import Self, is_typeddict from difflib import get_close_matches - +from langchain_google_vertexai.functions_utils import ( + _dict_to_gapic_schema, + _check_v2, +) logger = logging.getLogger(__name__) @@ -604,15 +607,13 @@ def _append_to_content( @overload def _parse_response_candidate( response_candidate: "Candidate", streaming: Literal[False] = False -) -> AIMessage: - ... +) -> AIMessage: ... @overload def _parse_response_candidate( response_candidate: "Candidate", streaming: Literal[True] -) -> AIMessageChunk: - ... +) -> AIMessageChunk: ... def _parse_response_candidate( @@ -2200,8 +2201,13 @@ class Explanation(BaseModel): if isinstance(schema, type) and is_basemodel_subclass(schema): if issubclass(schema, BaseModelV1): schema_json = schema.schema() + pydantic_version = "v1" else: schema_json = schema.model_json_schema() + pydantic_version = "v2" + schema_json = _dict_to_gapic_schema( + schema_json, pydantic_version=pydantic_version + ) parser = PydanticOutputParser(pydantic_object=schema) else: if is_typeddict(schema): @@ -2210,6 +2216,14 @@ class Explanation(BaseModel): schema_json = schema else: raise ValueError(f"Unsupported schema type {type(schema)}") + + pydantic_version_v2 = _check_v2(schema_json) + if pydantic_version_v2: + schema_json = _dict_to_gapic_schema( + schema_json, pydantic_version="v2" + ) + else: + schema_json = _dict_to_gapic_schema(schema_json) parser = JsonOutputParser() # Resolve refs in schema because they are not supported diff --git a/libs/vertexai/langchain_google_vertexai/functions_utils.py b/libs/vertexai/langchain_google_vertexai/functions_utils.py index c7f044f0c..da5beb434 100644 --- a/libs/vertexai/langchain_google_vertexai/functions_utils.py +++ b/libs/vertexai/langchain_google_vertexai/functions_utils.py @@ -230,11 +230,6 @@ def _format_pydantic_to_function_declaration( ) -def _format_dict_to_function_declaration( - tool: Union[FunctionDescription, Dict[str, Any]], -) -> gapic.FunctionDeclaration: - pydantic_version_v2 = False - # Ensure we send "anyOf" parameters through pydantic v2 schema parsing def _check_v2(parameters): properties = parameters.get("properties", {}).values() @@ -249,6 +244,12 @@ def _check_v2(parameters): return True return False + +def _format_dict_to_function_declaration( + tool: Union[FunctionDescription, Dict[str, Any]], +) -> gapic.FunctionDeclaration: + pydantic_version_v2 = False + if isinstance(tool, dict): pydantic_version_v2 = _check_v2(tool.get("parameters", {})) if pydantic_version_v2: From 98d8f0ba83605474fd1ffb16537965c962394943 Mon Sep 17 00:00:00 2001 From: yxia216 Date: Sat, 24 May 2025 16:33:46 -0400 Subject: [PATCH 2/2] update --- .../langchain_google_vertexai/chat_models.py | 14 ++++--- .../functions_utils.py | 37 ++++++++++++------- 2 files changed, 31 insertions(+), 20 deletions(-) diff --git a/libs/vertexai/langchain_google_vertexai/chat_models.py b/libs/vertexai/langchain_google_vertexai/chat_models.py index c1aefc45f..90c4d78f7 100644 --- a/libs/vertexai/langchain_google_vertexai/chat_models.py +++ b/libs/vertexai/langchain_google_vertexai/chat_models.py @@ -144,7 +144,7 @@ from typing_extensions import Self, is_typeddict from difflib import get_close_matches from langchain_google_vertexai.functions_utils import ( - _dict_to_gapic_schema, + _dict_to_gapic_schema_utils, _check_v2, ) @@ -607,13 +607,15 @@ def _append_to_content( @overload def _parse_response_candidate( response_candidate: "Candidate", streaming: Literal[False] = False -) -> AIMessage: ... +) -> AIMessage: + ... @overload def _parse_response_candidate( response_candidate: "Candidate", streaming: Literal[True] -) -> AIMessageChunk: ... +) -> AIMessageChunk: + ... def _parse_response_candidate( @@ -2205,7 +2207,7 @@ class Explanation(BaseModel): else: schema_json = schema.model_json_schema() pydantic_version = "v2" - schema_json = _dict_to_gapic_schema( + schema_json = _dict_to_gapic_schema_utils( schema_json, pydantic_version=pydantic_version ) parser = PydanticOutputParser(pydantic_object=schema) @@ -2219,11 +2221,11 @@ class Explanation(BaseModel): pydantic_version_v2 = _check_v2(schema_json) if pydantic_version_v2: - schema_json = _dict_to_gapic_schema( + schema_json = _dict_to_gapic_schema_utils( schema_json, pydantic_version="v2" ) else: - schema_json = _dict_to_gapic_schema(schema_json) + schema_json = _dict_to_gapic_schema_utils(schema_json) parser = JsonOutputParser() # Resolve refs in schema because they are not supported diff --git a/libs/vertexai/langchain_google_vertexai/functions_utils.py b/libs/vertexai/langchain_google_vertexai/functions_utils.py index da5beb434..da83a9f8a 100644 --- a/libs/vertexai/langchain_google_vertexai/functions_utils.py +++ b/libs/vertexai/langchain_google_vertexai/functions_utils.py @@ -165,9 +165,10 @@ def _format_json_schema_to_gapic( return converted_schema -def _dict_to_gapic_schema( +def _dict_to_gapic_schema_utils( schema: Dict[str, Any], pydantic_version: str = "v1" -) -> gapic.Schema: +) -> Dict[str, Any]: + """Convert the schema to make gemini understand.""" # Resolve refs in schema because $refs and $defs are not supported # by the Gemini API. dereferenced_schema = dereference_refs(schema) @@ -176,6 +177,13 @@ def _dict_to_gapic_schema( formatted_schema = _format_json_schema_to_gapic_v1(dereferenced_schema) else: formatted_schema = _format_json_schema_to_gapic(dereferenced_schema) + return formatted_schema + + +def _dict_to_gapic_schema( + schema: Dict[str, Any], pydantic_version: str = "v1" +) -> gapic.Schema: + formatted_schema = _dict_to_gapic_schema_utils(schema, pydantic_version) json_schema = json.dumps(formatted_schema) return gapic.Schema.from_json(json_schema) @@ -230,19 +238,20 @@ def _format_pydantic_to_function_declaration( ) - # Ensure we send "anyOf" parameters through pydantic v2 schema parsing - def _check_v2(parameters): - properties = parameters.get("properties", {}).values() - for property in properties: - if "anyOf" in property: +# Ensure we send "anyOf" parameters through pydantic v2 schema parsing +def _check_v2(parameters): + properties = parameters.get("properties", {}).values() + for property in properties: + if "anyOf" in property: + return True + if "parameters" in property: + if _check_v2(property["parameters"]): return True - if "parameters" in property: - if _check_v2(property["parameters"]): - return True - if "items" in property: - if _check_v2(property["items"]): - return True - return False + if "items" in property: + if _check_v2(property["items"]): + return True + + return False def _format_dict_to_function_declaration(