diff --git a/examples/01_basic_agent.py b/examples/01_basic_agent.py index e83bddc..b1111c5 100644 --- a/examples/01_basic_agent.py +++ b/examples/01_basic_agent.py @@ -88,10 +88,10 @@ async def main(): return # Example 2: Using Tokyo - print("\nExample 2: Using Tokyo") - print("-" * 50) - run = await get_capital_info.run(CityInput(city="Tokyo")) - print(run) + # print("\nExample 2: Using Tokyo") + # print("-" * 50) + # run = await get_capital_info.run(CityInput(city="Tokyo")) + # print(run) # Fetch and display completions for the Tokyo example await display_completions(run) diff --git a/pyproject.toml b/pyproject.toml index 02e28e1..f587ce7 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,6 @@ [tool.poetry] name = "workflowai" -version = "0.6.6" +version = "0.6.7" description = "Python SDK for WorkflowAI" authors = ["Guillaume Aquilina "] readme = "README.md" diff --git a/tests/fixtures/completions2.json b/tests/fixtures/completions2.json new file mode 100644 index 0000000..60d2c7a --- /dev/null +++ b/tests/fixtures/completions2.json @@ -0,0 +1,36 @@ +{ + "completions": [ + { + "messages": [ + { + "role": "assistant", + "content": [ + { + "type": "text", + "text": "\\nFind the capital city of the country where the input city is located.\\n\\nGuidelines:\\n1. First identify the country where the input city is located\\n2. Then provide the capital city of that country\\n3. Include an interesting historical or cultural fact about the capital\\n4. Be accurate and precise with geographical information\\n5. If the input city is itself the capital, still provide the information\\n\\n\\nInput will be provided in the user message using a JSON following the schema:\\n```json\\n{\\n \"description\": \"Input model for the city-to-capital agent.\",\\n \"properties\": {\\n \"city\": {\\n \"description\": \"The name of the city for which to find the country\\'s capital\",\\n \"examples\": [\\n \"Paris\",\\n \"New York\",\\n \"Tokyo\"\\n ],\\n \"type\": \"string\"\\n }\\n },\\n \"required\": [\\n \"city\"\\n ],\\n \"type\": \"object\"\\n}\\n```\\n\\nReturn a single JSON object enforcing the following schema:\\n```json\\n{\\n \"description\": \"Output model containing information about the capital city.\",\\n \"properties\": {\\n \"country\": {\\n \"description\": \"The country where the input city is located\",\\n \"examples\": [\\n \"France\",\\n \"United States\",\\n \"Japan\"\\n ],\\n \"type\": \"string\"\\n },\\n \"capital\": {\\n \"description\": \"The capital city of the country\",\\n \"examples\": [\\n \"Paris\",\\n \"Washington D.C.\",\\n \"Tokyo\"\\n ],\\n \"type\": \"string\"\\n },\\n \"fun_fact\": {\\n \"description\": \"An interesting fact about the capital city\",\\n \"examples\": [\\n \"Paris has been the capital of France since 508 CE\"\\n ],\\n \"type\": \"string\"\\n }\\n },\\n \"required\": [\\n \"capital\",\\n \"country\",\\n \"fun_fact\"\\n ],\\n \"type\": \"object\"\\n}\\n```" + } + ] + }, + { + "role": "user", + "content": [ + { + "type": "text", + "text": "Input is:\\n```json\\n{\\n \"city\": \"Paris\"\\n}\\n```" + } + ] + } + ], + "response": "{\\n \"country\": \"France\",\\n \"capital\": \"Paris\",\\n \"fun_fact\": \"Paris became the capital of France in 508 CE when Clovis I, the first King of the Franks, made the city his seat of power. It\\'s one of the oldest capital cities in Europe and was originally called Lutetia during the Roman period.\"\\n}", + "usage": { + "completion_token_count": 88, + "completion_cost_usd": 0.00132, + "prompt_token_count": 516, + "prompt_cost_usd": 0.0015480000000000001, + "model_context_window_size": 200000 + }, + "duration_seconds": 4.13, + "provider": "anthropic" + } + ] +} diff --git a/workflowai/core/client/agent_test.py b/workflowai/core/client/agent_test.py index b543803..1f8baf1 100644 --- a/workflowai/core/client/agent_test.py +++ b/workflowai/core/client/agent_test.py @@ -20,7 +20,7 @@ from workflowai.core.client.client import ( WorkflowAI, ) -from workflowai.core.domain.completion import Completion, CompletionUsage, Message +from workflowai.core.domain.completion import Completion, CompletionUsage, Message, TextContent from workflowai.core.domain.errors import MaxTurnsReachedError, WorkflowAIError from workflowai.core.domain.run import Run from workflowai.core.domain.tool_call import ToolCallRequest @@ -1131,6 +1131,24 @@ async def test_fetch_completions(self, agent: Agent[HelloTaskInput, HelloTaskOut ), ] + async def test_fetch_completions_content_list( + self, + agent: Agent[HelloTaskInput, HelloTaskOutput], + httpx_mock: HTTPXMock, + ): + """Test that fetch_completions correctly fetches and returns completions + when the message content is a list of objects""" + # Mock the HTTP response instead of the API client method + httpx_mock.add_response( + url="http://localhost:8000/v1/_/agents/123/runs/1/completions", + json=fixtures_json("completions2.json"), + ) + + completions = await agent.fetch_completions("1") + assert len(completions) == 1 + assert isinstance(completions[0].messages[0].content, list) + assert isinstance(completions[0].messages[0].content[0], TextContent) + class TestStream: async def test_stream(self, httpx_mock: HTTPXMock, agent: Agent[HelloTaskInput, HelloTaskOutput]): diff --git a/workflowai/core/domain/completion.py b/workflowai/core/domain/completion.py index 8572603..0c5a5ef 100644 --- a/workflowai/core/domain/completion.py +++ b/workflowai/core/domain/completion.py @@ -73,7 +73,7 @@ class Message(BaseModel): """A message in a completion.""" role: str = "" - content: Union[str, MessageContent] = Field(default="") + content: Union[str, list[MessageContent]] = Field(default="") class Completion(BaseModel): diff --git a/workflowai/core/domain/completion_test.py b/workflowai/core/domain/completion_test.py index fe951d6..433a5ca 100644 --- a/workflowai/core/domain/completion_test.py +++ b/workflowai/core/domain/completion_test.py @@ -17,70 +17,74 @@ def test_with_text_content(self): json_str = """ { "role": "assistant", - "content": { + "content": [{ "type": "text", "text": "This is a test message" - } + }] } """ message = Message.model_validate_json(json_str) assert message.role == "assistant" - assert isinstance(message.content, TextContent) - assert message.content.text == "This is a test message" + assert isinstance(message.content, list) + assert isinstance(message.content[0], TextContent) + assert message.content[0].text == "This is a test message" def test_with_document_content(self): # Test message with DocumentContent json_str = """ { "role": "user", - "content": { + "content": [{ "type": "document_url", "source": { "url": "https://example.com/doc.pdf" } - } + }] } """ message = Message.model_validate_json(json_str) assert message.role == "user" - assert isinstance(message.content, DocumentContent) - assert message.content.source.url == "https://example.com/doc.pdf" + assert isinstance(message.content, list) + assert isinstance(message.content[0], DocumentContent) + assert message.content[0].source.url == "https://example.com/doc.pdf" def test_with_image_content(self): # Test message with ImageContent json_str = """ { "role": "user", - "content": { + "content": [{ "type": "image_url", "image_url": { "url": "https://example.com/image.jpg" } - } + }] } """ message = Message.model_validate_json(json_str) assert message.role == "user" - assert isinstance(message.content, ImageContent) - assert message.content.image_url.url == "https://example.com/image.jpg" + assert isinstance(message.content, list) + assert isinstance(message.content[0], ImageContent) + assert message.content[0].image_url.url == "https://example.com/image.jpg" def test_with_audio_content(self): # Test message with AudioContent json_str = """ { "role": "user", - "content": { + "content": [{ "type": "audio_url", "audio_url": { "url": "https://example.com/audio.mp3" } - } + }] } """ message = Message.model_validate_json(json_str) assert message.role == "user" - assert isinstance(message.content, AudioContent) - assert message.content.audio_url.url == "https://example.com/audio.mp3" + assert isinstance(message.content, list) + assert isinstance(message.content[0], AudioContent) + assert message.content[0].audio_url.url == "https://example.com/audio.mp3" def test_empty_role(self): # Test message with empty role diff --git a/workflowai/core/domain/model.py b/workflowai/core/domain/model.py index 703ecce..2c521b1 100644 --- a/workflowai/core/domain/model.py +++ b/workflowai/core/domain/model.py @@ -139,5 +139,15 @@ class Model(str, Enum): DEEPSEEK_R1_2501 = "deepseek-r1-2501" DEEPSEEK_R1_2501_BASIC = "deepseek-r1-2501-basic" + # -------------------------------------------------------------------------- + # XAI Models + # -------------------------------------------------------------------------- + GROK_3_BETA = "grok-3-beta" + GROK_3_FAST_BETA = "grok-3-fast-beta" + GROK_3_MINI_BETA_LOW_REASONING_EFFORT = "grok-3-mini-beta-low" + GROK_3_MINI_BETA_HIGH_REASONING_EFFORT = "grok-3-mini-beta-high" + GROK_3_MINI_FAST_BETA_LOW_REASONING_EFFORT = "grok-3-mini-fast-beta-low" + GROK_3_MINI_FAST_BETA_HIGH_REASONING_EFFORT = "grok-3-mini-fast-beta-high" + ModelOrStr = Union[Model, str]