From 70b374fa29ca4649a97d893da1161b87ae789fd7 Mon Sep 17 00:00:00 2001 From: innocent-charles Date: Wed, 26 Aug 2026 03:52:21 +0300 Subject: [PATCH] Return chat responses as API JSON dicts by default. Drop raw=True/.raw for chat so callers can json.dumps the response directly and publish a cleaner client UX. Co-authored-by: Cursor --- CHANGELOG.md | 7 ++++ README.md | 18 ++++---- examples/README.md | 7 ++-- examples/getting_started.py | 42 ++++++------------- src/pawa_ai/_streaming.py | 73 +++++++++++++-------------------- src/pawa_ai/models/chat.py | 6 +-- src/pawa_ai/resources/agents.py | 33 ++++----------- src/pawa_ai/resources/chat.py | 41 ++++++------------ tests/test_client.py | 31 ++++---------- 9 files changed, 91 insertions(+), 167 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index db8822e..985a2bb 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -5,6 +5,13 @@ All notable changes to this project are documented in this file. The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.1.0/), and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html). +## [Unreleased] + +### Changed + +- `client.chat.create` and `client.agents.chat.create` always return the API JSON as a `dict` (no `raw=True` / `.raw`) +- `stream.collect()` returns a dict in the same shape as a non-streaming chat response + ## [0.2.0] - 2026-08-23 First public release on [PyPI](https://pypi.org/project/pawa-ai/). diff --git a/README.md b/README.md index 37b06cd..91c08ff 100644 --- a/README.md +++ b/README.md @@ -25,7 +25,7 @@ export PAWA_AI_API_KEY="your_api_key_here" Get your key from the [Builders Dashboard](https://builder.pawa-ai.com/dashboard?page=keys). -See **[examples/](examples/)** for a full walkthrough from `pip install` to typed responses. +See **[examples/](examples/)** for a full walkthrough from `pip install` to API responses. ### Chat @@ -45,8 +45,10 @@ response = client.chat.create( stream=False, ) -print(response.text) # typed ChatCompletion response -print(response.usage) # token usage when available +# Always a dict matching the API JSON +print(response["success"]) +print(response["data"]["request"][0]["message"]["content"]) +print(response["data"].get("usage")) ``` ### Streaming @@ -64,7 +66,7 @@ with client.chat.create( # Or collect the full response after streaming completion = stream.collect() - print(completion.text) + print(completion["data"]["request"][0]["message"]["content"]) ``` Async streaming: @@ -99,8 +101,6 @@ response = client.vectors.create( embeddings = response.embeddings ``` -Pass `raw=True` on any resource method to get the original JSON dict instead of typed models. - ### Retries with exponential backoff ```python @@ -133,7 +133,7 @@ async def main(): {"role": "user", "content": [{"type": "text", "text": "Habari yako?"}]} ], ) - print(response.text) + print(response["data"]["request"][0]["message"]["content"]) asyncio.run(main()) ``` @@ -156,12 +156,12 @@ asyncio.run(main()) ## Error handling ```python -from pawa_ai import PawaAI, AuthenticationError, RateLimitError, ChatCompletion +from pawa_ai import PawaAI, AuthenticationError, RateLimitError client = PawaAI() try: - completion: ChatCompletion = client.chat.create(model="pawa-v1-ember-20240924", messages=[...]) + response = client.chat.create(model="pawa-v1-ember-20240924", messages=[...]) except AuthenticationError as e: print(f"Auth failed: {e.message}") except RateLimitError as e: diff --git a/examples/README.md b/examples/README.md index 78383b8..0640db2 100644 --- a/examples/README.md +++ b/examples/README.md @@ -40,12 +40,11 @@ This script runs: | Example | What you get back | |---------|-------------------| -| Chat | `ChatCompletion` with `.text` and optional `.usage` | -| Streaming | Token deltas printed live, then full text via `.collect()` | +| Chat | API JSON `dict` (`success`, `message`, `data`) | +| Streaming | Token deltas printed live, then full dict via `.collect()` | | Embeddings | `EmbeddingResponse` with `.embeddings` (list of vectors) | | Models | `ModelList` with `.models` | | Error handling | Catches `AuthenticationError`, `RateLimitError`, etc. | -| Raw JSON | Plain `dict` when you pass `raw=True` | | Async chat | Same as chat, using `AsyncPawaAI` | ## 4. Minimal chat example @@ -65,7 +64,7 @@ response = client.chat.create( ], ) -print(response.text) +print(response["data"]["request"][0]["message"]["content"]) ``` ## 5. More capabilities diff --git a/examples/getting_started.py b/examples/getting_started.py index e724ad3..b92b8b9 100644 --- a/examples/getting_started.py +++ b/examples/getting_started.py @@ -29,7 +29,7 @@ def require_api_key() -> str: def example_chat() -> None: - """Basic chat completion — returns a typed ChatCompletion response.""" + """Basic chat completion — returns the API JSON as a dict.""" from pawa_ai import PawaAI client = PawaAI() @@ -56,12 +56,15 @@ def example_chat() -> None: stream=False, ) + reply = response["data"]["request"][0]["message"]["content"] + usage = response["data"].get("usage") or {} + print("=== Chat completion ===") - print(f"Success: {response.success}") - print(f"Model: {response.model}") - print(f"Reply: {response.text}") - if response.usage: - print(f"Tokens: in={response.usage.tokens_in}, out={response.usage.tokens_out}") + print(f"Success: {response['success']}") + print(f"Model: {response['data'].get('model')}") + print(f"Reply: {reply}") + if usage: + print(f"Tokens: in={usage.get('tokens_in')}, out={usage.get('tokens_out')}") print() @@ -88,7 +91,9 @@ def example_streaming() -> None: print(delta, end="", flush=True) print() - print(f"Collected: {stream.collect().text[:80]}...") + collected = stream.collect() + reply = collected["data"]["request"][0]["message"]["content"] + print(f"Collected: {reply[:80]}...") print() @@ -162,27 +167,7 @@ async def example_async_chat() -> None: ) print("=== Async chat ===") - print(f"Reply: {response.text}") - print() - - -def example_raw_response() -> None: - """Get the original JSON dict instead of typed models.""" - from pawa_ai import PawaAI - - client = PawaAI() - - payload = client.chat.create( - model="pawa-v1-ember-20240924", - messages=[ - {"role": "user", "content": [{"type": "text", "text": "Say hi in Swahili."}]} - ], - raw=True, - ) - - print("=== Raw JSON response ===") - print(f"Keys: {list(payload.keys())}") - print(f"Message field: {payload.get('message')}") + print(f"Reply: {response['data']['request'][0]['message']['content']}") print() @@ -196,7 +181,6 @@ def main() -> None: example_embeddings() example_list_models() example_error_handling() - example_raw_response() asyncio.run(example_async_chat()) diff --git a/src/pawa_ai/_streaming.py b/src/pawa_ai/_streaming.py index 8b1908e..91f9beb 100644 --- a/src/pawa_ai/_streaming.py +++ b/src/pawa_ai/_streaming.py @@ -1,9 +1,28 @@ from __future__ import annotations from collections.abc import AsyncIterator, Iterator +from typing import Any from pawa_ai._http import AsyncStream, Stream -from pawa_ai.models.chat import ChatCompletion, ChatStreamChunk +from pawa_ai.models.chat import ChatStreamChunk + + +def _collected_payload(text: str) -> dict[str, Any]: + return { + "success": True, + "message": "Stream collected", + "data": { + "request": [ + { + "finish_reason": "stop", + "message": {"role": "assistant", "content": text}, + } + ], + "created": "", + "model": "", + "object": "chat.request", + }, + } class ChatCompletionStream: @@ -26,8 +45,8 @@ def __exit__(self, *_: object) -> None: def chunks(self) -> Iterator[ChatStreamChunk]: """Yield typed stream chunks.""" - for raw in self._stream: - yield ChatStreamChunk.from_dict(raw) + for payload in self._stream: + yield ChatStreamChunk.from_dict(payload) def text_deltas(self) -> Iterator[str]: """Yield only the text delta from each chunk.""" @@ -39,26 +58,9 @@ def collect_text(self) -> str: """Collect all text deltas into a single string.""" return "".join(self.text_deltas()) - def collect(self) -> ChatCompletion: - """Build a :class:`ChatCompletion` from the full streamed text.""" - text = self.collect_text() - return ChatCompletion.from_dict( - { - "success": True, - "message": "Stream collected", - "data": { - "request": [ - { - "finish_reason": "stop", - "message": {"role": "assistant", "content": text}, - } - ], - "created": "", - "model": "", - "object": "chat.request", - }, - } - ) + def collect(self) -> dict[str, Any]: + """Build an API-shaped dict from the full streamed text.""" + return _collected_payload(self.collect_text()) def close(self) -> None: self._stream.close() @@ -74,8 +76,8 @@ def __aiter__(self) -> AsyncIterator[ChatStreamChunk]: return self.chunks() async def chunks(self) -> AsyncIterator[ChatStreamChunk]: - async for raw in self._stream: - yield ChatStreamChunk.from_dict(raw) + async for payload in self._stream: + yield ChatStreamChunk.from_dict(payload) async def text_deltas(self) -> AsyncIterator[str]: async for chunk in self.chunks(): @@ -88,25 +90,8 @@ async def collect_text(self) -> str: parts.append(delta) return "".join(parts) - async def collect(self) -> ChatCompletion: - text = await self.collect_text() - return ChatCompletion.from_dict( - { - "success": True, - "message": "Stream collected", - "data": { - "request": [ - { - "finish_reason": "stop", - "message": {"role": "assistant", "content": text}, - } - ], - "created": "", - "model": "", - "object": "chat.request", - }, - } - ) + async def collect(self) -> dict[str, Any]: + return _collected_payload(await self.collect_text()) async def close(self) -> None: await self._stream.close() diff --git a/src/pawa_ai/models/chat.py b/src/pawa_ai/models/chat.py index c5b69b3..67ecfc2 100644 --- a/src/pawa_ai/models/chat.py +++ b/src/pawa_ai/models/chat.py @@ -1,6 +1,6 @@ from __future__ import annotations -from dataclasses import dataclass, field +from dataclasses import dataclass from typing import Any @@ -72,7 +72,6 @@ class ChatCompletion: object: str choices: list[ChatChoice] usage: Usage | None = None - raw: dict[str, Any] = field(repr=False, default_factory=dict) @classmethod def from_dict(cls, payload: dict[str, Any]) -> ChatCompletion: @@ -87,7 +86,6 @@ def from_dict(cls, payload: dict[str, Any]) -> ChatCompletion: object=str(data.get("object", "")), choices=choices, usage=Usage.from_dict(data.get("usage")), - raw=payload, ) @property @@ -105,7 +103,6 @@ class ChatStreamChunk: message: str delta: str role: str | None = None - raw: dict[str, Any] = field(repr=False, default_factory=dict) @classmethod def from_dict(cls, payload: dict[str, Any]) -> ChatStreamChunk: @@ -118,5 +115,4 @@ def from_dict(cls, payload: dict[str, Any]) -> ChatStreamChunk: message=str(payload.get("message", "")), delta=delta, role=message_data.get("role"), - raw=payload, ) diff --git a/src/pawa_ai/resources/agents.py b/src/pawa_ai/resources/agents.py index 1501e54..54e29cb 100644 --- a/src/pawa_ai/resources/agents.py +++ b/src/pawa_ai/resources/agents.py @@ -2,9 +2,8 @@ from typing import TYPE_CHECKING, Any -from pawa_ai._http import Stream, raise_for_status +from pawa_ai._http import raise_for_status from pawa_ai._streaming import AsyncChatCompletionStream, ChatCompletionStream -from pawa_ai.models.chat import ChatCompletion if TYPE_CHECKING: from pawa_ai._client import AsyncPawaAI, PawaAI @@ -47,23 +46,17 @@ def __init__(self, client: PawaAI) -> None: def create( self, - *, - raw: bool = False, **params: Any, - ) -> ChatCompletion | ChatCompletionStream | Stream | dict[str, Any]: + ) -> dict[str, Any] | ChatCompletionStream: stream = bool(params.get("stream")) response = self._client._post("/agents/chat/request", json=params) if stream: - base_stream = Stream(response) - if raw: - return base_stream - return ChatCompletionStream(base_stream) + from pawa_ai._http import Stream + + return ChatCompletionStream(Stream(response)) raise_for_status(response) - payload = response.json() - if raw: - return payload - return ChatCompletion.from_dict(payload) + return response.json() class AsyncAgentsResource: @@ -103,22 +96,14 @@ def __init__(self, client: AsyncPawaAI) -> None: async def create( self, - *, - raw: bool = False, **params: Any, - ) -> ChatCompletion | AsyncChatCompletionStream | dict[str, Any]: + ) -> dict[str, Any] | AsyncChatCompletionStream: stream = bool(params.get("stream")) response = await self._client._post("/agents/chat/request", json=params) if stream: from pawa_ai._http import AsyncStream - base_stream = AsyncStream(response) - if raw: - return base_stream - return AsyncChatCompletionStream(base_stream) + return AsyncChatCompletionStream(AsyncStream(response)) raise_for_status(response) - payload = response.json() - if raw: - return payload - return ChatCompletion.from_dict(payload) + return response.json() diff --git a/src/pawa_ai/resources/chat.py b/src/pawa_ai/resources/chat.py index 09ea5dd..a05ea74 100644 --- a/src/pawa_ai/resources/chat.py +++ b/src/pawa_ai/resources/chat.py @@ -2,9 +2,8 @@ from typing import TYPE_CHECKING, Any -from pawa_ai._http import Stream, raise_for_status +from pawa_ai._http import raise_for_status from pawa_ai._streaming import AsyncChatCompletionStream, ChatCompletionStream -from pawa_ai.models.chat import ChatCompletion if TYPE_CHECKING: from pawa_ai._client import AsyncPawaAI, PawaAI @@ -17,30 +16,24 @@ def __init__(self, client: PawaAI) -> None: def create( self, - *, - raw: bool = False, **params: Any, - ) -> ChatCompletion | ChatCompletionStream | Stream | dict[str, Any]: + ) -> dict[str, Any] | ChatCompletionStream: """Send a chat completion request. - Set ``stream=True`` to receive a :class:`ChatCompletionStream`. - Pass ``raw=True`` to receive untyped dict/Stream responses. + Returns the API JSON as a ``dict``. Set ``stream=True`` to receive a + :class:`ChatCompletionStream`. """ stream = bool(params.get("stream")) response = self._client._post("/chat/request", json=params) if stream: - base_stream = Stream(response) - if raw: - return base_stream - return ChatCompletionStream(base_stream) + from pawa_ai._http import Stream + + return ChatCompletionStream(Stream(response)) raise_for_status(response) - payload = response.json() - if raw: - return payload - return ChatCompletion.from_dict(payload) + return response.json() - def completions(self, **params: Any) -> ChatCompletion | ChatCompletionStream | dict[str, Any]: + def completions(self, **params: Any) -> dict[str, Any] | ChatCompletionStream: """Alias for :meth:`create`.""" return self.create(**params) @@ -51,27 +44,19 @@ def __init__(self, client: AsyncPawaAI) -> None: async def create( self, - *, - raw: bool = False, **params: ChatRequest | Any, - ) -> ChatCompletion | AsyncChatCompletionStream | dict[str, Any]: + ) -> dict[str, Any] | AsyncChatCompletionStream: stream = bool(params.get("stream")) response = await self._client._post("/chat/request", json=params) if stream: from pawa_ai._http import AsyncStream - base_stream = AsyncStream(response) - if raw: - return base_stream - return AsyncChatCompletionStream(base_stream) + return AsyncChatCompletionStream(AsyncStream(response)) raise_for_status(response) - payload = response.json() - if raw: - return payload - return ChatCompletion.from_dict(payload) + return response.json() async def completions( self, **params: ChatRequest | Any - ) -> ChatCompletion | AsyncChatCompletionStream | dict[str, Any]: + ) -> dict[str, Any] | AsyncChatCompletionStream: return await self.create(**params) diff --git a/tests/test_client.py b/tests/test_client.py index 261bfab..15a2160 100644 --- a/tests/test_client.py +++ b/tests/test_client.py @@ -6,14 +6,12 @@ from pawa_ai import ( AsyncPawaAI, AuthenticationError, - ChatCompletion, ChatCompletionStream, PawaAI, RetryConfig, ) from pawa_ai._http import parse_stream_line, raise_for_status from pawa_ai._streaming import AsyncChatCompletionStream -from pawa_ai.models.chat import ChatCompletion as ChatCompletionModel from pawa_ai.models.chat import ChatStreamChunk @@ -38,7 +36,7 @@ def test_client_builds_auth_header(): assert client._http_client is not None -def test_chat_create_returns_typed_model(): +def test_chat_create_returns_dict(): payload = { "success": True, "message": "Chat request processed successfully", @@ -69,24 +67,10 @@ def handler(request: httpx.Request) -> httpx.Response: messages=[{"role": "user", "content": [{"type": "text", "text": "Hi"}]}], stream=False, ) - assert isinstance(response, ChatCompletion) - assert response.success is True - assert response.text == "Hello there" - assert response.usage is not None - assert response.usage.tokens_in == 10 - - -def test_chat_create_raw_response(): - payload = {"success": True, "message": "ok", "data": {}} - transport = httpx.MockTransport(lambda request: httpx.Response(200, json=payload)) - client = PawaAI(api_key="test-key", http_client=httpx.Client(transport=transport)) - - response = client.chat.create( - model="pawa-v1-ember-20240924", - messages=[{"role": "user", "content": [{"type": "text", "text": "Hi"}]}], - raw=True, - ) + assert isinstance(response, dict) assert response["success"] is True + assert response["data"]["request"][0]["message"]["content"] == "Hello there" + assert response["data"]["usage"]["tokens_in"] == 10 def test_raise_for_status_maps_errors(): @@ -122,8 +106,8 @@ def handler(request: httpx.Request) -> httpx.Response: messages=[{"role": "user", "content": [{"type": "text", "text": "Hi"}]}], stream=True, ).collect() - assert isinstance(collected, ChatCompletionModel) - assert collected.text == "Hello" + assert isinstance(collected, dict) + assert collected["data"]["request"][0]["message"]["content"] == "Hello" @pytest.mark.asyncio @@ -153,7 +137,7 @@ async def handler(request: httpx.Request) -> httpx.Response: model="pawa-v1-ember-20240924", messages=[{"role": "user", "content": [{"type": "text", "text": "Hi"}]}], ) - assert response.text == "A" + assert response["data"]["request"][0]["message"]["content"] == "A" @pytest.mark.asyncio @@ -198,7 +182,6 @@ def handler(request: httpx.Request) -> httpx.Response: response = client.chat.create( model="pawa-v1-ember-20240924", messages=[{"role": "user", "content": [{"type": "text", "text": "Hi"}]}], - raw=True, ) assert attempts["count"] == 2