From 145818110e6e7b914dbbcfb99e327984fb32e7bd Mon Sep 17 00:00:00 2001 From: dannyjameswilliams Date: Fri, 11 Sep 2026 14:10:27 +0100 Subject: [PATCH 1/7] Add image generation support --- test/test_imports.py | 6 ++ uv.lock | 2 +- weaviate_agents/classes/__init__.py | 3 + weaviate_agents/classes/media.py | 31 ++++++ weaviate_agents/query/query_agent.py | 135 +++++++++++++++++++++++++-- 5 files changed, 168 insertions(+), 9 deletions(-) create mode 100644 weaviate_agents/classes/media.py diff --git a/test/test_imports.py b/test/test_imports.py index 52e44aa..b62e790 100644 --- a/test/test_imports.py +++ b/test/test_imports.py @@ -59,6 +59,7 @@ def test_class_exports(): DependentOperationStep, FilterAndOr, GeoPropertyFilter, + ImageOptions, IntegerArrayPropertyFilter, IntegerPropertyAggregation, IntegerPropertyFilter, @@ -76,6 +77,7 @@ def test_class_exports(): PersonalizedQueryResponse, ProgressDetails, ProgressMessage, + QAImage, QueryAgentCollectionConfig, QueryAgentResponse, QueryResult, @@ -133,11 +135,13 @@ def test_class_exports(): DatePropertyAggregation, DatePropertyFilter, DependentOperationStep, + ImageOptions, IntegerArrayPropertyFilter, IntegerPropertyAggregation, IntegerPropertyFilter, IsNullPropertyFilter, GeoPropertyFilter, + QAImage, NumericMetrics, Operations, OperationStep, @@ -193,6 +197,8 @@ def test_class_exports(): "QueryWithCollection", "Source", "ChatMessage", + "QAImage", + "ImageOptions", "ComparisonOperator", "IntegerPropertyFilter", "TextPropertyFilter", diff --git a/uv.lock b/uv.lock index 10f3b58..19f2837 100644 --- a/uv.lock +++ b/uv.lock @@ -1044,7 +1044,7 @@ wheels = [ [[package]] name = "weaviate-agents" -version = "1.5.0" +version = "1.7.1" source = { editable = "." } dependencies = [ { name = "httpx-sse" }, diff --git a/weaviate_agents/classes/__init__.py b/weaviate_agents/classes/__init__.py index 5ae9ac7..145f042 100644 --- a/weaviate_agents/classes/__init__.py +++ b/weaviate_agents/classes/__init__.py @@ -1,4 +1,5 @@ from .core import Usage +from .media import ImageOptions, QAImage from .personalization import ( Persona, PersonaInteraction, @@ -115,6 +116,8 @@ "IsNullPropertyFilter", "SearchModeResponseBase", "ChatMessage", + "QAImage", + "ImageOptions", "AskModeResponse", "ResearchModeResponse", "ModelUnitUsage", diff --git a/weaviate_agents/classes/media.py b/weaviate_agents/classes/media.py new file mode 100644 index 0000000..0810833 --- /dev/null +++ b/weaviate_agents/classes/media.py @@ -0,0 +1,31 @@ +from dataclasses import dataclass +from typing import Optional + +from pydantic import BaseModel, ConfigDict, GetJsonSchemaHandler +from pydantic.json_schema import JsonSchemaValue, SkipJsonSchema +from pydantic_core import core_schema +from typing_extensions import Literal + +IMAGE_KEYWORD = "X-query-agent-image" + +ImageShape = Literal["square", "landscape", "portrait"] + + +class QAImage(BaseModel): + model_config = ConfigDict(json_schema_extra={IMAGE_KEYWORD: True}) + + image_prompt: str + base64: SkipJsonSchema[str] # hidden from the LLM's schema; server fills it + + +@dataclass(frozen=True) +class ImageOptions: + shape: Optional[ImageShape] = None # unset falls back to the backend default + + def __get_pydantic_json_schema__( + self, core_schema_: core_schema.CoreSchema, handler: GetJsonSchemaHandler + ) -> JsonSchemaValue: + schema = handler(core_schema_) + if self.shape is not None: + schema["X-image-shape"] = self.shape + return schema diff --git a/weaviate_agents/query/query_agent.py b/weaviate_agents/query/query_agent.py index b9fdc53..539485a 100644 --- a/weaviate_agents/query/query_agent.py +++ b/weaviate_agents/query/query_agent.py @@ -21,6 +21,7 @@ from weaviate.client import WeaviateAsyncClient, WeaviateClient from weaviate_agents.base import ClientType, _BaseAgent +from weaviate_agents.classes.media import IMAGE_KEYWORD from weaviate_agents.query.classes import ( AskModeResponse, ParsedAskModeResponse, @@ -44,6 +45,8 @@ # type through to `ParsedAskModeResponse[MyModel].final_answer_parsed`. M = TypeVar("M", bound=BaseModel) +_MEDIA_REQUEST_TIMEOUT = 180 # seconds + class _BaseQueryAgent(Generic[ClientType], _BaseAgent[ClientType], ABC): def __init__( @@ -60,9 +63,21 @@ def __init__( self._system_prompt = system_prompt self._timeout = 60 if timeout is None else timeout + self._timeout_explicit = timeout is not None self.agent_url = f"{self._agents_host}/agent" self.query_url = f"{self._agents_host}/query" + def _request_timeout( + self, output_format: Union[None, dict[str, Any], type[BaseModel]] = None + ) -> int: + if ( + not self._timeout_explicit + and output_format is not None + and _schema_contains_media(output_format) + ): + return _MEDIA_REQUEST_TIMEOUT + return self._timeout + def _prepare_request_body( self, query: Union[str, list[ChatMessage]], @@ -824,7 +839,8 @@ def ask( The LLM will conform to the output format specified. The `final_answer_parsed` output field in the response will also be of the type specified. When passing a `dict`, the dictionary must conform to the Draft 2020-12 JSON Schema specification. - Defaults to `str`. + To generate an image, include a :class:`~weaviate_agents.classes.media.QAImage` field within a Pydantic model (as `output_format`). + The base64 of the image will be returned in the corresponding field of `final_answer_parsed` as string. Returns: An instance of :class:`~weaviate_agents.query.classes.response.AskModeResponse` (or [:class:`~weaviate_agents.query.classes.response.ParsedAskModeResponse`] @@ -855,6 +871,21 @@ def ask( >>> result = agent.ask("What contracts were signed by Jane Doe in 2024? What were they about?", output_format=AnswerWithSources) >>> print(type(result.final_answer_parsed)) + + >>> from weaviate_agents.classes.media import QAImage + >>> from base64 import b64decode + >>> class AnswerWithImage(BaseModel): + ... answer: str + ... image: QAImage = Field( + ... description="An advertisement for the product.", + ... ) + >>> + >>> agent = QueryAgent( + ... client=client, + ... collections=["ECommerce"], + ... ) + >>> result = agent.ask("What is the best selling product?", output_format=AnswerWithImage) + >>> image_bytes = b64decode(result.final_answer_parsed.image.base64) # base64 encoded image """ request_body = self._prepare_request_body( query=query, @@ -867,7 +898,7 @@ def ask( self.query_url + "/ask", headers=self._headers, json=request_body, - timeout=self._timeout, + timeout=self._request_timeout(output_format=request_body["output_format"]), ) if response.is_error: @@ -1136,6 +1167,10 @@ def ask_stream( Whilst streaming, the :class:`~weaviate_agents.query.classes.response.StreamedTokens` will return delta text tokens on the final answer as it is being constructed as raw string tokens, not a JSON object. When the final answer is complete, the :class:`~weaviate_agents.query.classes.response.ParsedAskModeResponse` will be returned if ``include_final_state`` is ``True``. + To generate an image, include a :class:`~weaviate_agents.classes.media.QAImage` field within a Pydantic model (as `output_format`). + The base64 of the image will be returned in the corresponding field of `final_answer_parsed` as string. + If a image is being generated, the streamed tokens of the final answer will not include the `"base64"` key, it is added afterwards + and will be present in the final result. Returns: A generator of the response stream. @@ -1153,7 +1188,7 @@ def ask_stream( ... collections=["FinancialContracts"], ... ) >>> for result in agent.ask_stream("What are the terms of the contract signed by John Smith in May 2025?"): - ... if isinstance(result, AskModeResponse): + ... if isinstance(result, AskModeResponse): # this will also find ParsedAskModeResponse ... result.display() ... elif isinstance(result, StreamedTokens): ... print(result.delta, end='', flush=True) @@ -1178,12 +1213,32 @@ def ask_stream( ... "What contracts were signed by Jane Doe in 2024? What were they about?", ... output_format=AnswerWithSources ... ): - ... if isinstance(result, AskModeResponse): # this will also find ParsedAskModeResponse + ... if isinstance(result, AskModeResponse): ... result.display() ... elif isinstance(result, StreamedTokens): ... print(result.delta, end='', flush=True) ... elif isinstance(result, ProgressMessage): ... print(result.message) + + >>> from weaviate_agents.classes.media import QAImage + >>> from base64 import b64decode + >>> class AnswerWithImage(BaseModel): + ... answer: str + ... image: QAImage = Field( + ... description="An advertisement for the product.", + ... ) + >>> + >>> agent = QueryAgent( + ... client=client, + ... collections=["ECommerce"], + ... ) + >>> for result in agent.ask_stream("What is the best selling product?", output_format=AnswerWithImage): + ... if isinstance(result, AskModeResponse): # the final result; `base64` is only populated here + ... image_bytes = b64decode(result.final_answer_parsed.image.base64) + ... elif isinstance(result, StreamedTokens): + ... print(result.delta, end='', flush=True) + ... elif isinstance(result, ProgressMessage): + ... print(result.message) """ request_body = self._prepare_request_body( query=query, @@ -1200,7 +1255,9 @@ def ask_stream( url=self.query_url + "/stream_ask", json=request_body, headers=self._headers, - timeout=self._timeout, + timeout=self._request_timeout( + output_format=request_body["output_format"] + ), ) as events: if events.response.is_error: events.response.read() @@ -1673,7 +1730,8 @@ async def ask( The LLM will conform to the output format specified. The `final_answer_parsed` output field in the response will also be of the type specified. When passing a `dict`, the dictionary must conform to the Draft 2020-12 JSON Schema specification. - Defaults to `str`. + To generate an image, include a :class:`~weaviate_agents.classes.media.QAImage` field within a Pydantic model (as `output_format`). + The base64 of the image will be returned in the corresponding field of `final_answer_parsed` as string. Returns: An instance of :class:`~weaviate_agents.query.classes.response.AskModeResponse` (or [:class:`~weaviate_agents.query.classes.response.ParsedAskModeResponse`] if ``output_format`` is not ``None``) which contains the final answer, sources, @@ -1704,6 +1762,21 @@ async def ask( >>> result = await agent.ask("What contracts were signed by Jane Doe in 2024? What were they about?", output_format=AnswerWithSources) >>> print(type(result.final_answer_parsed)) + + >>> from weaviate_agents.classes.media import QAImage + >>> from base64 import b64decode + >>> class AnswerWithImage(BaseModel): + ... answer: str + ... image: QAImage = Field( + ... description="An advertisement for the product.", + ... ) + >>> + >>> agent = AsyncQueryAgent( + ... client=client, + ... collections=["ECommerce"], + ... ) + >>> result = await agent.ask("What is the best selling product?", output_format=AnswerWithImage) + >>> image_bytes = b64decode(result.final_answer_parsed.image.base64) # base64 encoded image """ request_body = self._prepare_request_body( query=query, @@ -1717,7 +1790,9 @@ async def ask( self.query_url + "/ask", headers=self._headers, json=request_body, - timeout=self._timeout, + timeout=self._request_timeout( + output_format=request_body["output_format"] + ), ) if response.is_error: @@ -1989,6 +2064,10 @@ async def ask_stream( Whilst streaming, the :class:`~weaviate_agents.query.classes.response.StreamedTokens` will return delta text tokens on the final answer as it is being constructed as raw string tokens, not a JSON object. When the final answer is complete, the :class:`~weaviate_agents.query.classes.response.ParsedAskModeResponse` will be returned if ``include_final_state`` is ``True``. + To generate an image, include a :class:`~weaviate_agents.classes.media.QAImage` field within a Pydantic model (as `output_format`). + The base64 of the image will be returned in the corresponding field of `final_answer_parsed` as string. + If a image is being generated, the streamed tokens of the final answer will not include the `"base64"` key, it is added afterwards + and will be present in the final result. Returns: A generator of the response stream. @@ -2037,6 +2116,26 @@ async def ask_stream( ... print(result.delta, end='', flush=True) ... elif isinstance(result, ProgressMessage): ... print(result.message) + + >>> from weaviate_agents.classes.media import QAImage + >>> from base64 import b64decode + >>> class AnswerWithImage(BaseModel): + ... answer: str + ... image: QAImage = Field( + ... description="An advertisement for the product.", + ... ) + >>> + >>> agent = AsyncQueryAgent( + ... client=client, + ... collections=["ECommerce"], + ... ) + >>> async for result in agent.ask_stream("What is the best selling product?", output_format=AnswerWithImage): + ... if isinstance(result, AskModeResponse): # the final result; `base64` is only populated here + ... image_bytes = b64decode(result.final_answer_parsed.image.base64) + ... elif isinstance(result, StreamedTokens): + ... print(result.delta, end='', flush=True) + ... elif isinstance(result, ProgressMessage): + ... print(result.message) """ request_body = self._prepare_request_body( query=query, @@ -2053,7 +2152,9 @@ async def ask_stream( url=self.query_url + "/stream_ask", json=request_body, headers=self._headers, - timeout=self._timeout, + timeout=self._request_timeout( + output_format=request_body["output_format"] + ), ) as events: if events.response.is_error: await events.response.aread() @@ -2508,3 +2609,21 @@ def _parse_ask_result( return ParsedAskModeResponse[dict](**response) return AskModeResponse(**response) + + +def _schema_contains_media( + schema: Union[dict[str, Any], list[Any], type[BaseModel]], +) -> bool: + """Return True if the schema (or any nested node) requests generated media. + + Used to increase timeouts when any nested node is tagged with the media keyword. + """ + if isinstance(schema, type) and issubclass(schema, BaseModel): + schema = schema.model_json_schema() + if isinstance(schema, dict): + if IMAGE_KEYWORD in schema: + return True + return any(_schema_contains_media(value) for value in schema.values()) + if isinstance(schema, list): + return any(_schema_contains_media(item) for item in schema) + return False From c148865d3a1d193f0d342f02387a87e42f7bd168 Mon Sep 17 00:00:00 2001 From: dannyjameswilliams Date: Fri, 25 Sep 2026 10:33:52 +0100 Subject: [PATCH 2/7] allow annotated at root of output_format --- test/query_agent/test_query_model.py | 42 ++++++++++++++++++++++++++++ weaviate_agents/query/query_agent.py | 31 ++++++++++++++++++-- 2 files changed, 70 insertions(+), 3 deletions(-) diff --git a/test/query_agent/test_query_model.py b/test/query_agent/test_query_model.py index 587b40e..204323e 100644 --- a/test/query_agent/test_query_model.py +++ b/test/query_agent/test_query_model.py @@ -1,11 +1,13 @@ import json from contextlib import asynccontextmanager, contextmanager +from typing import Annotated, List import httpx import pytest from httpx_sse import ServerSentEvent from pydantic import ValidationError +from weaviate_agents.classes.media import ImageOptions, QAImage from weaviate_agents.classes.query import ( AskModeResponse, ProgressMessage, @@ -1355,6 +1357,46 @@ def fake_post_with_capture(url, headers=None, json=None, timeout=None): assert captured["json"]["query"] == {"messages": chat_messages} +def test_ask_with_annotated_image_output_format(monkeypatch): + captured = {} + image = {"image_prompt": "a red shoe", "base64": "QUJD"} + response = {**FAKE_ASK_SUCCESS_JSON, "final_answer": json.dumps(image)} + + def fake_post_with_capture(url, headers=None, json=None, timeout=None): + captured["json"] = json + return FakeResponse(200, response) + + monkeypatch.setattr(httpx, "post", fake_post_with_capture) + dummy_client = DummyClient() + agent = QueryAgent( + dummy_client, ["test_collection"], agents_host="http://dummy-agent" + ) + agent._connection = dummy_client + agent._headers = dummy_client.additional_headers + + with pytest.warns(UserWarning): + result = agent.ask( + "draw a shoe", + output_format=Annotated[QAImage, ImageOptions(shape="square")], + ) + + # the schema is sent with its shape, and the answer parses back into a QAImage + assert captured["json"]["output_format"]["X-image-shape"] == "square" + assert result.final_answer_parsed == QAImage(**image) + + +def test_ask_rejects_unsupported_output_format(): + dummy_client = DummyClient() + agent = QueryAgent( + dummy_client, ["test_collection"], agents_host="http://dummy-agent" + ) + agent._connection = dummy_client + + # a list root used to be dropped silently, returning a plain-text answer + with pytest.raises(TypeError): + agent.ask("draw shoes", output_format=List[QAImage]) + + def test_ask_failure(monkeypatch): monkeypatch.setattr(httpx, "post", fake_post_failure) dummy_client = DummyClient() diff --git a/weaviate_agents/query/query_agent.py b/weaviate_agents/query/query_agent.py index 539485a..0a408f2 100644 --- a/weaviate_agents/query/query_agent.py +++ b/weaviate_agents/query/query_agent.py @@ -2,6 +2,7 @@ from abc import ABC, abstractmethod from json import JSONDecodeError, loads from typing import ( + Annotated, Any, AsyncGenerator, Coroutine, @@ -11,17 +12,19 @@ Optional, TypeVar, Union, + get_args, + get_origin, overload, ) import httpx from httpx_sse import ServerSentEvent, aconnect_sse, connect_sse -from pydantic import BaseModel +from pydantic import BaseModel, TypeAdapter from typing_extensions import deprecated from weaviate.client import WeaviateAsyncClient, WeaviateClient from weaviate_agents.base import ClientType, _BaseAgent -from weaviate_agents.classes.media import IMAGE_KEYWORD +from weaviate_agents.classes.media import IMAGE_KEYWORD, QAImage from weaviate_agents.query.classes import ( AskModeResponse, ParsedAskModeResponse, @@ -107,10 +110,18 @@ def _prepare_request_body( if isinstance(output_format, type) and issubclass(output_format, BaseModel): output_format_json = output_format.model_json_schema() + elif _is_annotated_image(output_format): + # the schema comes from the Annotated form so it keeps the ImageOptions + output_format_json = TypeAdapter(output_format).json_schema() elif isinstance(output_format, dict): output_format_json = output_format - else: + elif output_format is None: output_format_json = None + else: + raise TypeError( + "output_format must be a BaseModel subclass, a dict JSON schema, " + f"QAImage, or Annotated[QAImage, ImageOptions(...)], got {output_format!r}" + ) output = { "query": query_request, @@ -2598,6 +2609,12 @@ def _parse_ask_result( ) return ParsedAskModeResponse[BaseModel](**response) + elif _is_annotated_image(output_format): + response["final_answer_parsed"] = TypeAdapter(output_format).validate_json( + response["final_answer"] + ) + return ParsedAskModeResponse[BaseModel](**response) + elif isinstance(output_format, dict): try: response["final_answer_parsed"] = loads(response["final_answer"]) @@ -2611,6 +2628,14 @@ def _parse_ask_result( return AskModeResponse(**response) +def _is_annotated_image(output_format: Any) -> bool: + """Whether output_format is an Annotated[QAImage, ...] root, e.g. carrying ImageOptions.""" + if get_origin(output_format) is not Annotated: + return False + base = get_args(output_format)[0] + return isinstance(base, type) and issubclass(base, QAImage) + + def _schema_contains_media( schema: Union[dict[str, Any], list[Any], type[BaseModel]], ) -> bool: From 0c1eed5b637c2e9576c919cfde6756eb9ed0bdde Mon Sep 17 00:00:00 2001 From: dannyjameswilliams Date: Mon, 5 Oct 2026 10:32:57 +0100 Subject: [PATCH 3/7] reduce timeout -> 60s --- weaviate_agents/query/query_agent.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/weaviate_agents/query/query_agent.py b/weaviate_agents/query/query_agent.py index 0a408f2..fbd5172 100644 --- a/weaviate_agents/query/query_agent.py +++ b/weaviate_agents/query/query_agent.py @@ -48,7 +48,7 @@ # type through to `ParsedAskModeResponse[MyModel].final_answer_parsed`. M = TypeVar("M", bound=BaseModel) -_MEDIA_REQUEST_TIMEOUT = 180 # seconds +_MEDIA_REQUEST_TIMEOUT = 60 # seconds class _BaseQueryAgent(Generic[ClientType], _BaseAgent[ClientType], ABC): From 55abdb17f934a7fa22dd7a2443b00c23ebf0c9f5 Mon Sep 17 00:00:00 2001 From: dannyjameswilliams Date: Tue, 6 Oct 2026 11:15:21 +0100 Subject: [PATCH 4/7] rename QAImage -> GeneratedImage --- test/query_agent/test_query_model.py | 11 +++++----- test/test_imports.py | 6 +++--- uv.lock | 2 +- weaviate_agents/classes/__init__.py | 4 ++-- weaviate_agents/classes/media.py | 2 +- weaviate_agents/query/query_agent.py | 32 ++++++++++++++-------------- 6 files changed, 28 insertions(+), 29 deletions(-) diff --git a/test/query_agent/test_query_model.py b/test/query_agent/test_query_model.py index f03b3a7..bc573e2 100644 --- a/test/query_agent/test_query_model.py +++ b/test/query_agent/test_query_model.py @@ -7,7 +7,7 @@ from httpx_sse import ServerSentEvent from pydantic import ValidationError -from weaviate_agents.classes.media import ImageOptions, QAImage +from weaviate_agents.classes.media import GeneratedImage, ImageOptions from weaviate_agents.classes.query import ( AskModeResponse, ProgressMessage, @@ -1452,12 +1452,12 @@ def fake_post_with_capture(url, headers=None, json=None, timeout=None): with pytest.warns(UserWarning): result = agent.ask( "draw a shoe", - output_format=Annotated[QAImage, ImageOptions(shape="square")], + output_format=Annotated[GeneratedImage, ImageOptions(shape="square")], ) - # the schema is sent with its shape, and the answer parses back into a QAImage + # the schema is sent with its shape, and the answer parses back into a GeneratedImage assert captured["json"]["output_format"]["X-image-shape"] == "square" - assert result.final_answer_parsed == QAImage(**image) + assert result.final_answer_parsed == GeneratedImage(**image) def test_ask_rejects_unsupported_output_format(): @@ -1467,9 +1467,8 @@ def test_ask_rejects_unsupported_output_format(): ) agent._connection = dummy_client - # a list root used to be dropped silently, returning a plain-text answer with pytest.raises(TypeError): - agent.ask("draw shoes", output_format=List[QAImage]) + agent.ask("draw shoes", output_format=List[GeneratedImage]) def test_ask_failure(monkeypatch): diff --git a/test/test_imports.py b/test/test_imports.py index b62e790..000d093 100644 --- a/test/test_imports.py +++ b/test/test_imports.py @@ -58,6 +58,7 @@ def test_class_exports(): DatePropertyFilter, DependentOperationStep, FilterAndOr, + GeneratedImage, GeoPropertyFilter, ImageOptions, IntegerArrayPropertyFilter, @@ -77,7 +78,6 @@ def test_class_exports(): PersonalizedQueryResponse, ProgressDetails, ProgressMessage, - QAImage, QueryAgentCollectionConfig, QueryAgentResponse, QueryResult, @@ -141,7 +141,7 @@ def test_class_exports(): IntegerPropertyFilter, IsNullPropertyFilter, GeoPropertyFilter, - QAImage, + GeneratedImage, NumericMetrics, Operations, OperationStep, @@ -197,7 +197,7 @@ def test_class_exports(): "QueryWithCollection", "Source", "ChatMessage", - "QAImage", + "GeneratedImage", "ImageOptions", "ComparisonOperator", "IntegerPropertyFilter", diff --git a/uv.lock b/uv.lock index 19f2837..d2c082c 100644 --- a/uv.lock +++ b/uv.lock @@ -1044,7 +1044,7 @@ wheels = [ [[package]] name = "weaviate-agents" -version = "1.7.1" +version = "1.8.0" source = { editable = "." } dependencies = [ { name = "httpx-sse" }, diff --git a/weaviate_agents/classes/__init__.py b/weaviate_agents/classes/__init__.py index 145f042..f0c347b 100644 --- a/weaviate_agents/classes/__init__.py +++ b/weaviate_agents/classes/__init__.py @@ -1,5 +1,5 @@ from .core import Usage -from .media import ImageOptions, QAImage +from .media import GeneratedImage, ImageOptions from .personalization import ( Persona, PersonaInteraction, @@ -116,7 +116,7 @@ "IsNullPropertyFilter", "SearchModeResponseBase", "ChatMessage", - "QAImage", + "GeneratedImage", "ImageOptions", "AskModeResponse", "ResearchModeResponse", diff --git a/weaviate_agents/classes/media.py b/weaviate_agents/classes/media.py index 0810833..7d580ea 100644 --- a/weaviate_agents/classes/media.py +++ b/weaviate_agents/classes/media.py @@ -11,7 +11,7 @@ ImageShape = Literal["square", "landscape", "portrait"] -class QAImage(BaseModel): +class GeneratedImage(BaseModel): model_config = ConfigDict(json_schema_extra={IMAGE_KEYWORD: True}) image_prompt: str diff --git a/weaviate_agents/query/query_agent.py b/weaviate_agents/query/query_agent.py index d599cf4..81cef40 100644 --- a/weaviate_agents/query/query_agent.py +++ b/weaviate_agents/query/query_agent.py @@ -24,7 +24,7 @@ from weaviate.client import WeaviateAsyncClient, WeaviateClient from weaviate_agents.base import ClientType, _BaseAgent -from weaviate_agents.classes.media import IMAGE_KEYWORD, QAImage +from weaviate_agents.classes.media import IMAGE_KEYWORD, GeneratedImage from weaviate_agents.query.classes import ( AskModeResponse, ParsedAskModeResponse, @@ -120,7 +120,7 @@ def _prepare_request_body( else: raise TypeError( "output_format must be a BaseModel subclass, a dict JSON schema, " - f"QAImage, or Annotated[QAImage, ImageOptions(...)], got {output_format!r}" + f"GeneratedImage, or Annotated[GeneratedImage, ImageOptions(...)], got {output_format!r}" ) output = { @@ -850,7 +850,7 @@ def ask( The LLM will conform to the output format specified. The `final_answer_parsed` output field in the response will also be of the type specified. When passing a `dict`, the dictionary must conform to the Draft 2020-12 JSON Schema specification. - To generate an image, include a :class:`~weaviate_agents.classes.media.QAImage` field within a Pydantic model (as `output_format`). + To generate an image, include a :class:`~weaviate_agents.classes.media.GeneratedImage` field within a Pydantic model (as `output_format`). The base64 of the image will be returned in the corresponding field of `final_answer_parsed` as string. Returns: @@ -883,11 +883,11 @@ def ask( >>> print(type(result.final_answer_parsed)) - >>> from weaviate_agents.classes.media import QAImage + >>> from weaviate_agents.classes.media import GeneratedImage >>> from base64 import b64decode >>> class AnswerWithImage(BaseModel): ... answer: str - ... image: QAImage = Field( + ... image: GeneratedImage = Field( ... description="An advertisement for the product.", ... ) >>> @@ -1178,7 +1178,7 @@ def ask_stream( Whilst streaming, the :class:`~weaviate_agents.query.classes.response.StreamedTokens` will return delta text tokens on the final answer as it is being constructed as raw string tokens, not a JSON object. When the final answer is complete, the :class:`~weaviate_agents.query.classes.response.ParsedAskModeResponse` will be returned if ``include_final_state`` is ``True``. - To generate an image, include a :class:`~weaviate_agents.classes.media.QAImage` field within a Pydantic model (as `output_format`). + To generate an image, include a :class:`~weaviate_agents.classes.media.GeneratedImage` field within a Pydantic model (as `output_format`). The base64 of the image will be returned in the corresponding field of `final_answer_parsed` as string. If a image is being generated, the streamed tokens of the final answer will not include the `"base64"` key, it is added afterwards and will be present in the final result. @@ -1231,11 +1231,11 @@ def ask_stream( ... elif isinstance(result, ProgressMessage): ... print(result.message) - >>> from weaviate_agents.classes.media import QAImage + >>> from weaviate_agents.classes.media import GeneratedImage >>> from base64 import b64decode >>> class AnswerWithImage(BaseModel): ... answer: str - ... image: QAImage = Field( + ... image: GeneratedImage = Field( ... description="An advertisement for the product.", ... ) >>> @@ -1746,7 +1746,7 @@ async def ask( The LLM will conform to the output format specified. The `final_answer_parsed` output field in the response will also be of the type specified. When passing a `dict`, the dictionary must conform to the Draft 2020-12 JSON Schema specification. - To generate an image, include a :class:`~weaviate_agents.classes.media.QAImage` field within a Pydantic model (as `output_format`). + To generate an image, include a :class:`~weaviate_agents.classes.media.GeneratedImage` field within a Pydantic model (as `output_format`). The base64 of the image will be returned in the corresponding field of `final_answer_parsed` as string. Returns: @@ -1779,11 +1779,11 @@ async def ask( >>> print(type(result.final_answer_parsed)) - >>> from weaviate_agents.classes.media import QAImage + >>> from weaviate_agents.classes.media import GeneratedImage >>> from base64 import b64decode >>> class AnswerWithImage(BaseModel): ... answer: str - ... image: QAImage = Field( + ... image: GeneratedImage = Field( ... description="An advertisement for the product.", ... ) >>> @@ -2080,7 +2080,7 @@ async def ask_stream( Whilst streaming, the :class:`~weaviate_agents.query.classes.response.StreamedTokens` will return delta text tokens on the final answer as it is being constructed as raw string tokens, not a JSON object. When the final answer is complete, the :class:`~weaviate_agents.query.classes.response.ParsedAskModeResponse` will be returned if ``include_final_state`` is ``True``. - To generate an image, include a :class:`~weaviate_agents.classes.media.QAImage` field within a Pydantic model (as `output_format`). + To generate an image, include a :class:`~weaviate_agents.classes.media.GeneratedImage` field within a Pydantic model (as `output_format`). The base64 of the image will be returned in the corresponding field of `final_answer_parsed` as string. If a image is being generated, the streamed tokens of the final answer will not include the `"base64"` key, it is added afterwards and will be present in the final result. @@ -2133,11 +2133,11 @@ async def ask_stream( ... elif isinstance(result, ProgressMessage): ... print(result.message) - >>> from weaviate_agents.classes.media import QAImage + >>> from weaviate_agents.classes.media import GeneratedImage >>> from base64 import b64decode >>> class AnswerWithImage(BaseModel): ... answer: str - ... image: QAImage = Field( + ... image: GeneratedImage = Field( ... description="An advertisement for the product.", ... ) >>> @@ -2639,11 +2639,11 @@ def _parse_ask_result( def _is_annotated_image(output_format: Any) -> bool: - """Whether output_format is an Annotated[QAImage, ...] root, e.g. carrying ImageOptions.""" + """Whether output_format is an Annotated[GeneratedImage, ...] root, e.g. carrying ImageOptions.""" if get_origin(output_format) is not Annotated: return False base = get_args(output_format)[0] - return isinstance(base, type) and issubclass(base, QAImage) + return isinstance(base, type) and issubclass(base, GeneratedImage) def _schema_contains_media( From 4c5242193049e8abf42834d0e05bd949df4bca2d Mon Sep 17 00:00:00 2001 From: dannyjameswilliams Date: Tue, 6 Oct 2026 12:06:26 +0100 Subject: [PATCH 5/7] rename ImageOptions -> GeneratedImageOptions --- test/query_agent/test_query_model.py | 6 ++++-- test/test_imports.py | 6 +++--- weaviate_agents/classes/__init__.py | 4 ++-- weaviate_agents/classes/media.py | 2 +- weaviate_agents/query/query_agent.py | 6 +++--- 5 files changed, 13 insertions(+), 11 deletions(-) diff --git a/test/query_agent/test_query_model.py b/test/query_agent/test_query_model.py index bc573e2..f7a1f4d 100644 --- a/test/query_agent/test_query_model.py +++ b/test/query_agent/test_query_model.py @@ -7,7 +7,7 @@ from httpx_sse import ServerSentEvent from pydantic import ValidationError -from weaviate_agents.classes.media import GeneratedImage, ImageOptions +from weaviate_agents.classes.media import GeneratedImage, GeneratedImageOptions from weaviate_agents.classes.query import ( AskModeResponse, ProgressMessage, @@ -1452,7 +1452,9 @@ def fake_post_with_capture(url, headers=None, json=None, timeout=None): with pytest.warns(UserWarning): result = agent.ask( "draw a shoe", - output_format=Annotated[GeneratedImage, ImageOptions(shape="square")], + output_format=Annotated[ + GeneratedImage, GeneratedImageOptions(shape="square") + ], ) # the schema is sent with its shape, and the answer parses back into a GeneratedImage diff --git a/test/test_imports.py b/test/test_imports.py index 000d093..eba0063 100644 --- a/test/test_imports.py +++ b/test/test_imports.py @@ -59,8 +59,8 @@ def test_class_exports(): DependentOperationStep, FilterAndOr, GeneratedImage, + GeneratedImageOptions, GeoPropertyFilter, - ImageOptions, IntegerArrayPropertyFilter, IntegerPropertyAggregation, IntegerPropertyFilter, @@ -135,7 +135,7 @@ def test_class_exports(): DatePropertyAggregation, DatePropertyFilter, DependentOperationStep, - ImageOptions, + GeneratedImageOptions, IntegerArrayPropertyFilter, IntegerPropertyAggregation, IntegerPropertyFilter, @@ -198,7 +198,7 @@ def test_class_exports(): "Source", "ChatMessage", "GeneratedImage", - "ImageOptions", + "GeneratedImageOptions", "ComparisonOperator", "IntegerPropertyFilter", "TextPropertyFilter", diff --git a/weaviate_agents/classes/__init__.py b/weaviate_agents/classes/__init__.py index f0c347b..9f2a39c 100644 --- a/weaviate_agents/classes/__init__.py +++ b/weaviate_agents/classes/__init__.py @@ -1,5 +1,5 @@ from .core import Usage -from .media import GeneratedImage, ImageOptions +from .media import GeneratedImage, GeneratedImageOptions from .personalization import ( Persona, PersonaInteraction, @@ -117,7 +117,7 @@ "SearchModeResponseBase", "ChatMessage", "GeneratedImage", - "ImageOptions", + "GeneratedImageOptions", "AskModeResponse", "ResearchModeResponse", "ModelUnitUsage", diff --git a/weaviate_agents/classes/media.py b/weaviate_agents/classes/media.py index 7d580ea..8346c49 100644 --- a/weaviate_agents/classes/media.py +++ b/weaviate_agents/classes/media.py @@ -19,7 +19,7 @@ class GeneratedImage(BaseModel): @dataclass(frozen=True) -class ImageOptions: +class GeneratedImageOptions: shape: Optional[ImageShape] = None # unset falls back to the backend default def __get_pydantic_json_schema__( diff --git a/weaviate_agents/query/query_agent.py b/weaviate_agents/query/query_agent.py index 81cef40..828ec50 100644 --- a/weaviate_agents/query/query_agent.py +++ b/weaviate_agents/query/query_agent.py @@ -111,7 +111,7 @@ def _prepare_request_body( if isinstance(output_format, type) and issubclass(output_format, BaseModel): output_format_json = output_format.model_json_schema() elif _is_annotated_image(output_format): - # the schema comes from the Annotated form so it keeps the ImageOptions + # the schema comes from the Annotated form so it keeps the GeneratedImageOptions output_format_json = TypeAdapter(output_format).json_schema() elif isinstance(output_format, dict): output_format_json = output_format @@ -120,7 +120,7 @@ def _prepare_request_body( else: raise TypeError( "output_format must be a BaseModel subclass, a dict JSON schema, " - f"GeneratedImage, or Annotated[GeneratedImage, ImageOptions(...)], got {output_format!r}" + f"GeneratedImage, or Annotated[GeneratedImage, GeneratedImageOptions(...)], got {output_format!r}" ) output = { @@ -2639,7 +2639,7 @@ def _parse_ask_result( def _is_annotated_image(output_format: Any) -> bool: - """Whether output_format is an Annotated[GeneratedImage, ...] root, e.g. carrying ImageOptions.""" + """Whether output_format is an Annotated[GeneratedImage, ...] root, e.g. carrying GeneratedImageOptions.""" if get_origin(output_format) is not Annotated: return False base = get_args(output_format)[0] From 8d11571222aa8542a4fc635fe590d90864a74f61 Mon Sep 17 00:00:00 2001 From: dannyjameswilliams Date: Tue, 6 Oct 2026 14:07:38 +0100 Subject: [PATCH 6/7] remove timeout increase --- weaviate_agents/query/query_agent.py | 28 ++++------------------------ 1 file changed, 4 insertions(+), 24 deletions(-) diff --git a/weaviate_agents/query/query_agent.py b/weaviate_agents/query/query_agent.py index 828ec50..11aa9c2 100644 --- a/weaviate_agents/query/query_agent.py +++ b/weaviate_agents/query/query_agent.py @@ -48,8 +48,6 @@ # type through to `ParsedAskModeResponse[MyModel].final_answer_parsed`. M = TypeVar("M", bound=BaseModel) -_MEDIA_REQUEST_TIMEOUT = 60 # seconds - class _BaseQueryAgent(Generic[ClientType], _BaseAgent[ClientType], ABC): def __init__( @@ -66,21 +64,9 @@ def __init__( self._system_prompt = system_prompt self._timeout = 60 if timeout is None else timeout - self._timeout_explicit = timeout is not None self.agent_url = f"{self._agents_host}/agent" self.query_url = f"{self._agents_host}/query" - def _request_timeout( - self, output_format: Union[None, dict[str, Any], type[BaseModel]] = None - ) -> int: - if ( - not self._timeout_explicit - and output_format is not None - and _schema_contains_media(output_format) - ): - return _MEDIA_REQUEST_TIMEOUT - return self._timeout - def _prepare_request_body( self, query: Union[str, list[ChatMessage]], @@ -909,7 +895,7 @@ def ask( self.query_url + "/ask", headers=self._headers, json=request_body, - timeout=self._request_timeout(output_format=request_body["output_format"]), + timeout=self._timeout, ) if response.is_error: @@ -1266,9 +1252,7 @@ def ask_stream( url=self.query_url + "/stream_ask", json=request_body, headers=self._headers, - timeout=self._request_timeout( - output_format=request_body["output_format"] - ), + timeout=self._timeout, ) as events: if events.response.is_error: events.response.read() @@ -1806,9 +1790,7 @@ async def ask( self.query_url + "/ask", headers=self._headers, json=request_body, - timeout=self._request_timeout( - output_format=request_body["output_format"] - ), + timeout=self._timeout, ) if response.is_error: @@ -2168,9 +2150,7 @@ async def ask_stream( url=self.query_url + "/stream_ask", json=request_body, headers=self._headers, - timeout=self._request_timeout( - output_format=request_body["output_format"] - ), + timeout=self._timeout, ) as events: if events.response.is_error: await events.response.aread() From c01b8f5dc95e10ddd44dfb054ddf02acf730553c Mon Sep 17 00:00:00 2001 From: dannyjameswilliams Date: Thu, 15 Oct 2026 12:07:02 +0100 Subject: [PATCH 7/7] remove refernece to root Annotated field --- weaviate_agents/query/query_agent.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/weaviate_agents/query/query_agent.py b/weaviate_agents/query/query_agent.py index 11aa9c2..ff98c3b 100644 --- a/weaviate_agents/query/query_agent.py +++ b/weaviate_agents/query/query_agent.py @@ -105,8 +105,7 @@ def _prepare_request_body( output_format_json = None else: raise TypeError( - "output_format must be a BaseModel subclass, a dict JSON schema, " - f"GeneratedImage, or Annotated[GeneratedImage, GeneratedImageOptions(...)], got {output_format!r}" + f"output_format must be a BaseModel subclass or a dict JSON schema, got {output_format!r}" ) output = {