Skip to content

Commit

Permalink
fix: Vertex AI: handle "default" in Field
Browse files Browse the repository at this point in the history
fixes jxnl#837
  • Loading branch information
cmitsakis committed Jul 30, 2024
1 parent 8407377 commit 2db79be
Show file tree
Hide file tree
Showing 2 changed files with 12 additions and 4 deletions.
10 changes: 9 additions & 1 deletion instructor/client_vertexai.py
Original file line number Diff line number Diff line change
Expand Up @@ -12,10 +12,18 @@
def _create_gemini_json_schema(model: BaseModel):
schema = model.model_json_schema()
schema_without_refs: dict[str, Any] = jsonref.replace_refs(schema) # type: ignore
try:
required_params = schema_without_refs["required"]
except KeyError:
# `schema_without_refs` has no key "required" because all parameters have "default" value set
required_params = []
# find parameters that have "default" value set, and add them to `required_params`
params_with_default = [param for param, description in schema_without_refs["properties"].items() if "default" in description]
required_params += params_with_default
gemini_schema: dict[Any, Any] = {
"type": schema_without_refs["type"],
"properties": schema_without_refs["properties"],
"required": schema_without_refs["required"],
"required": required_params,
}
return gemini_schema

Expand Down
6 changes: 3 additions & 3 deletions tests/llm/test_vertexai/test_modes.py
Original file line number Diff line number Diff line change
Expand Up @@ -14,7 +14,7 @@ class Item(BaseModel):

class Order(BaseModel):
items: list[Item] = Field(..., default_factory=list)
customer: str
customer: str = Field(default="")


@pytest.mark.parametrize("model, mode", product(models, modes))
Expand Down Expand Up @@ -49,13 +49,13 @@ def test_nested(model, mode):
class Book(BaseModel):
title: str
author: str
genre: str
genre: str = Field(default="")
isbn: str


class LibraryRecord(BaseModel):
books: list[Book] = Field(..., default_factory=list)
visitor: str
visitor: str = Field(default="")
library_id: str


Expand Down

0 comments on commit 2db79be

Please sign in to comment.