Merge pull request #8 from Sartify/fix/chat-return-dict

Return chat responses as API JSON dicts by default.

Author: innocent-charles

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/).

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:

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

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())
 

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()

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,
         )

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()