Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

vertexai[patch]: remove allOf from function schema #96

Open
wants to merge 2 commits into
base: main
Choose a base branch
from
Open

vertexai[patch]: remove allOf from function schema #96

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
07ca7a3
vertexai[patch]: remove allOf from function schema
baskaryan Mar 26, 2024
e3addb9
fmt
baskaryan Mar 26, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Jump to
Jump to file
Failed to load files.
Diff view
Diff view
31 changes: 29 additions & 2 deletions libs/vertexai/langchain_google_vertexai/functions_utils.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -94,6 +94,31 @@ class Config:
extra = "allow"


def _replace_all_ofs(schema: Dict[str, Any]) -> Dict[str, Any]:
"""Assumes dereferenced schema."""
new_schema: Dict[str, Any] = {}
for k, v in schema.items():
if isinstance(v, dict) and "allOf" in v:
if isinstance(v["allOf"], list) and len(v["allOf"]) == 1:
obj = dict(v["allOf"][0])
obj["title"] = v.get("title", obj.get("title", ""))
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

title field is now not supported - we are dropping it

obj["description"] = " ".join(
(v.get("description", ""), obj.get("description", ""))
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I'm not sure about joining descriptions.

This code:

class Node(BaseModel):
    """This is a node"""
    id: str
    type: str

class Relationship(BaseModel):
    source: Node
    target: Node = Field(..., description="This is a target node")

will produce following definition for target property:

{
  "target": {
    "title": "Target",
    "description": "This is a target node This is a node",
    "type": "object",
    "properties": {
      "id": {
        "title": "Id",
        "type": "string"
      },
      "type": {
        "title": "Type",
        "type": "string"
      }
    },
    "required": ["id", "type"]
  }
}

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

yea not ideal but is there a better generic solution? both descriptions could have important (and distinct) info

Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I would consider coalescing i. e. field description should take precedence

class Relationship(BaseModel):
    source: Node
    target: Node = Field(..., description="This is a target node")

over model description:

class Node(BaseModel):
    """This is a node"""
    id: str
    type: str

And could you please describe reasoning and behaviour in the function doc string?

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

i can get behind that, will update!

Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@lkuligin what is your take on that?

)
new_schema[k] = obj
else:
raise ValueError(
f"allOf expected to be a singleton list. Received {v['allOf']}"
)
elif isinstance(v, dict):
new_schema[k] = _replace_all_ofs(v)
elif isinstance(v, list) and v and isinstance(v[0], dict):
new_schema[k] = [_replace_all_ofs(s) for s in v]
else:
new_schema[k] = v
return new_schema


def _get_parameters_from_schema(schema: Dict[str, Any]) -> Dict[str, Any]:
"""Given a schema, format the parameters key to match VertexAI
expected input.
Expand All @@ -105,8 +130,10 @@ def _get_parameters_from_schema(schema: Dict[str, Any]) -> Dict[str, Any]:
Dictionary with the formatted parameters.
"""

dereferenced_schema = dereference_refs(schema)
model = ParametersSchema.parse_obj(dereferenced_schema)
schema = dereference_refs(schema)
# TODO: Remove if vertexai api supports allOf elements in the future.
schema = _replace_all_ofs(schema)
model = ParametersSchema.parse_obj(schema)

return model.dict(exclude_unset=True)

Expand Down
45 changes: 45 additions & 0 deletions libs/vertexai/tests/unit_tests/test_function_utils.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -6,6 +6,7 @@
from langchain_google_vertexai.functions_utils import (
_format_tool_to_vertex_function,
_get_parameters_from_schema,
_replace_all_ofs,
)


Expand Down Expand Up @@ -77,3 +78,47 @@ class B(BaseModel):
assert len(result["properties"]["b1"]["properties"]) == 1

assert "anyOf" in result["properties"]["b3"]


def test__replace_all_ofs() -> None:
schema = {
"title": "Relationship",
"type": "object",
"properties": {
"target": {
"title": "Target",
"description": "The target of the relationship.",
"allOf": [
{
"title": "Node",
"type": "object",
"description": "An entity in a graph.",
"properties": {
"id": {"title": "Id", "type": "string"},
},
"required": ["id"],
}
],
}
},
"required": ["target"],
}
expected = {
"title": "Relationship",
"type": "object",
"properties": {
"target": {
"title": "Target",
"type": "object",
"description": "The target of the relationship. An entity in a graph.",
"properties": {
"id": {"title": "Id", "type": "string"},
},
"required": ["id"],
}
},
"required": ["target"],
}
actual = _replace_all_ofs(schema)
assert actual == expected
assert _replace_all_ofs(expected) == expected