diff --git a/google/genai/tests/models/test_generate_content.py b/google/genai/tests/models/test_generate_content.py index bd85ba3..5dfd2f1 100644 --- a/google/genai/tests/models/test_generate_content.py +++ b/google/genai/tests/models/test_generate_content.py @@ -404,6 +404,60 @@ def test_json_schema(client): assert isinstance(response.parsed, dict) +def test_json_schema_with_streaming(client): + + response = client.models.generate_content_stream( + model='gemini-2.0-flash-exp', + contents='Give me information of the United States.', + config={ + 'response_mime_type': 'application/json', + 'response_schema': { + 'properties': { + 'name': {'type': 'STRING'}, + 'population': {'type': 'INTEGER'}, + 'capital': {'type': 'STRING'}, + 'continent': {'type': 'STRING'}, + 'gdp': {'type': 'INTEGER'}, + 'official_language': {'type': 'STRING'}, + 'total_area_sq_mi': {'type': 'INTEGER'}, + }, + 'type': 'OBJECT', + }, + }, + ) + + for r in response: + parts = r.candidates[0].content.parts + for p in parts: + print(p.text) + + +def test_pydantic_schema_with_streaming(client): + + class CountryInfo(BaseModel): + name: str + population: int + capital: str + continent: str + gdp: int + official_language: str + total_area_sq_mi: int + + response = client.models.generate_content_stream( + model='gemini-2.0-flash-exp', + contents='Give me information of the United States.', + config={ + 'response_mime_type': 'application/json', + 'response_schema': CountryInfo + }, + ) + + for r in response: + parts = r.candidates[0].content.parts + for p in parts: + print(p.text) + + def test_function(client): def get_weather(city: str) -> str: return f'The weather in {city} is sunny and 100 degrees.' diff --git a/google/genai/types.py b/google/genai/types.py index ebf154c..2594dcb 100644 --- a/google/genai/types.py +++ b/google/genai/types.py @@ -2453,7 +2453,10 @@ class GenerateContentResponse(_common.BaseModel): default=None, description="""Usage metadata about the response(s).""" ) automatic_function_calling_history: Optional[list[Content]] = None - parsed: Union[pydantic.BaseModel, dict] = None + parsed: Union[pydantic.BaseModel, dict] = Field( + default=None, + description="""Parsed response if response_schema is provided. Not available for streaming.""", + ) @property def text(self) -> Optional[str]: @@ -2503,12 +2506,21 @@ def _from_response( response_schema, pydantic.BaseModel ): # Pydantic schema. - result.parsed = response_schema.model_validate_json(result.text) + try: + result.parsed = response_schema.model_validate_json(result.text) + # may not be a valid json per stream response + except pydantic.ValidationError: + pass + elif isinstance(response_schema, dict) or isinstance( response_schema, pydantic.BaseModel ): # JSON schema. - result.parsed = json.loads(result.text) + try: + result.parsed = json.loads(result.text) + # may not be a valid json per stream response + except json.decoder.JSONDecodeError: + pass return result