diff --git a/test/query_agent/test_query_model.py b/test/query_agent/test_query_model.py index 67420a9..f7a1f4d 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 GeneratedImage, GeneratedImageOptions from weaviate_agents.classes.query import ( AskModeResponse, ProgressMessage, @@ -1430,6 +1432,47 @@ 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[ + GeneratedImage, GeneratedImageOptions(shape="square") + ], + ) + + # 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 == GeneratedImage(**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 + + with pytest.raises(TypeError): + agent.ask("draw shoes", output_format=List[GeneratedImage]) + + def test_ask_failure(monkeypatch): monkeypatch.setattr(httpx, "post", fake_post_failure) dummy_client = DummyClient() diff --git a/test/test_imports.py b/test/test_imports.py index 52e44aa..eba0063 100644 --- a/test/test_imports.py +++ b/test/test_imports.py @@ -58,6 +58,8 @@ def test_class_exports(): DatePropertyFilter, DependentOperationStep, FilterAndOr, + GeneratedImage, + GeneratedImageOptions, GeoPropertyFilter, IntegerArrayPropertyFilter, IntegerPropertyAggregation, @@ -133,11 +135,13 @@ def test_class_exports(): DatePropertyAggregation, DatePropertyFilter, DependentOperationStep, + GeneratedImageOptions, IntegerArrayPropertyFilter, IntegerPropertyAggregation, IntegerPropertyFilter, IsNullPropertyFilter, GeoPropertyFilter, + GeneratedImage, NumericMetrics, Operations, OperationStep, @@ -193,6 +197,8 @@ def test_class_exports(): "QueryWithCollection", "Source", "ChatMessage", + "GeneratedImage", + "GeneratedImageOptions", "ComparisonOperator", "IntegerPropertyFilter", "TextPropertyFilter", diff --git a/uv.lock b/uv.lock index 10f3b58..d2c082c 100644 --- a/uv.lock +++ b/uv.lock @@ -1044,7 +1044,7 @@ wheels = [ [[package]] name = "weaviate-agents" -version = "1.5.0" +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 5ae9ac7..9f2a39c 100644 --- a/weaviate_agents/classes/__init__.py +++ b/weaviate_agents/classes/__init__.py @@ -1,4 +1,5 @@ from .core import Usage +from .media import GeneratedImage, GeneratedImageOptions from .personalization import ( Persona, PersonaInteraction, @@ -115,6 +116,8 @@ "IsNullPropertyFilter", "SearchModeResponseBase", "ChatMessage", + "GeneratedImage", + "GeneratedImageOptions", "AskModeResponse", "ResearchModeResponse", "ModelUnitUsage", diff --git a/weaviate_agents/classes/media.py b/weaviate_agents/classes/media.py new file mode 100644 index 0000000..8346c49 --- /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 GeneratedImage(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 GeneratedImageOptions: + 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 3a0df10..ff98c3b 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,16 +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, GeneratedImage from weaviate_agents.query.classes import ( AskModeResponse, ParsedAskModeResponse, @@ -92,10 +96,17 @@ 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 GeneratedImageOptions + 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( + f"output_format must be a BaseModel subclass or a dict JSON schema, got {output_format!r}" + ) output = { "query": query_request, @@ -824,7 +835,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.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: An instance of :class:`~weaviate_agents.query.classes.response.AskModeResponse` (or [:class:`~weaviate_agents.query.classes.response.ParsedAskModeResponse`] @@ -855,6 +867,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 GeneratedImage + >>> from base64 import b64decode + >>> class AnswerWithImage(BaseModel): + ... answer: str + ... image: GeneratedImage = 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, @@ -1136,6 +1163,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.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. Returns: A generator of the response stream. @@ -1153,7 +1184,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 +1209,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 GeneratedImage + >>> from base64 import b64decode + >>> class AnswerWithImage(BaseModel): + ... answer: str + ... image: GeneratedImage = 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, @@ -1678,7 +1729,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.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: 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, @@ -1709,6 +1761,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 GeneratedImage + >>> from base64 import b64decode + >>> class AnswerWithImage(BaseModel): + ... answer: str + ... image: GeneratedImage = 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, @@ -1994,6 +2061,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.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. Returns: A generator of the response stream. @@ -2042,6 +2113,26 @@ async def ask_stream( ... print(result.delta, end='', flush=True) ... elif isinstance(result, ProgressMessage): ... print(result.message) + + >>> from weaviate_agents.classes.media import GeneratedImage + >>> from base64 import b64decode + >>> class AnswerWithImage(BaseModel): + ... answer: str + ... image: GeneratedImage = 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, @@ -2507,6 +2598,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"]) @@ -2518,3 +2615,29 @@ def _parse_ask_result( return ParsedAskModeResponse[dict](**response) return AskModeResponse(**response) + + +def _is_annotated_image(output_format: Any) -> bool: + """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] + return isinstance(base, type) and issubclass(base, GeneratedImage) + + +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