add stream

Author: CLEMENTINATOR

README.md

@@ -71,6 +71,34 @@ if response.choices[0].message.get("tool_calls"):
     print(response.choices[0].message["tool_calls"])
 ```
 
+### Streaming
+
+```python
+for chunk in edgee.stream(
+    model="gpt-4o",
+    input="Tell me a story",
+):
+    if chunk.choices[0].delta.content:
+        print(chunk.choices[0].delta.content, end="", flush=True)
+```
+
+### Streaming with Messages
+
+```python
+for chunk in edgee.stream(
+    model="gpt-4o",
+    input={
+        "messages": [
+            {"role": "system", "content": "You are a helpful assistant."},
+            {"role": "user", "content": "Hello!"},
+        ],
+    },
+):
+    delta = chunk.choices[0].delta
+    if delta.content:
+        print(delta.content, end="", flush=True)
+```
+
 ## Response
 
 ```python
@@ -92,4 +120,24 @@ class Usage:
     total_tokens: int
 ```
 
+### Streaming Response
+
+```python
+@dataclass
+class StreamChunk:
+    choices: list[StreamChoice]
+
+@dataclass
+class StreamChoice:
+    index: int
+    delta: StreamDelta
+    finish_reason: str | None
+
+@dataclass
+class StreamDelta:
+    role: str | None  # Only present in first chunk
+    content: str | None
+    tool_calls: list[dict] | None
+```
+
 To learn more about this SDK, please refer to the [dedicated documentation](https://www.edgee.cloud/docs/sdk/python).

edgee/__init__.py

@@ -67,6 +67,25 @@ class SendResponse:
     usage: Usage | None = None
 
 
+@dataclass
+class StreamDelta:
+    role: str | None = None
+    content: str | None = None
+    tool_calls: list[dict] | None = None
+
+
+@dataclass
+class StreamChoice:
+    index: int
+    delta: StreamDelta
+    finish_reason: str | None = None
+
+
+@dataclass
+class StreamChunk:
+    choices: list[StreamChoice]
+
+
 @dataclass
 class EdgeeConfig:
     api_key: str | None = None
@@ -159,3 +178,87 @@ def send(
             )
 
         return SendResponse(choices=choices, usage=usage)
+
+    def stream(
+        self,
+        model: str,
+        input: str | InputObject | dict,
+    ):
+        """Stream a completion request from the Edgee AI Gateway.
+
+        Yields StreamChunk objects as they arrive from the API.
+        """
+
+        if isinstance(input, str):
+            messages = [{"role": "user", "content": input}]
+            tools = None
+            tool_choice = None
+        elif isinstance(input, InputObject):
+            messages = input.messages
+            tools = input.tools
+            tool_choice = input.tool_choice
+        else:
+            messages = input.get("messages", [])
+            tools = input.get("tools")
+            tool_choice = input.get("tool_choice")
+
+        body: dict = {"model": model, "messages": messages, "stream": True}
+        if tools:
+            body["tools"] = tools
+        if tool_choice:
+            body["tool_choice"] = tool_choice
+
+        request = Request(
+            f"{self.base_url}{API_ENDPOINT}",
+            data=json.dumps(body).encode("utf-8"),
+            headers={
+                "Content-Type": "application/json",
+                "Authorization": f"Bearer {self.api_key}",
+            },
+            method="POST",
+        )
+
+        try:
+            with urlopen(request) as response:
+                # Read and parse SSE stream
+                buffer = ""
+                for line in response:
+                    decoded_line = line.decode("utf-8")
+
+                    if decoded_line.strip() == "":
+                        continue
+
+                    if decoded_line.startswith("data: "):
+                        data_str = decoded_line[6:].strip()
+
+                        # Check for stream end signal
+                        if data_str == "[DONE]":
+                            break
+
+                        try:
+                            data = json.loads(data_str)
+
+                            # Parse choices
+                            choices = []
+                            for c in data.get("choices", []):
+                                delta_data = c.get("delta", {})
+                                delta = StreamDelta(
+                                    role=delta_data.get("role"),
+                                    content=delta_data.get("content"),
+                                    tool_calls=delta_data.get("tool_calls"),
+                                )
+                                choice = StreamChoice(
+                                    index=c["index"],
+                                    delta=delta,
+                                    finish_reason=c.get("finish_reason"),
+                                )
+                                choices.append(choice)
+
+                            yield StreamChunk(choices=choices)
+                        except json.JSONDecodeError:
+                            # Skip malformed JSON
+                            continue
+
+        except HTTPError as e:
+            error_body = e.read().decode("utf-8")
+            raise RuntimeError(f"API error {e.code}: {error_body}") from e

example/test.py

@@ -13,7 +13,7 @@
 # Test 1: Simple string input
 print("Test 1: Simple string input")
 response1 = edgee.send(
-    model="gpt-4o",
+    model="mistral/mistral-small-latest",
     input="What is the capital of France?",
 )
 print(f"Content: {response1.choices[0].message['content']}")
@@ -23,7 +23,7 @@
 # Test 2: Full input object with messages
 print("Test 2: Full input object with messages")
 response2 = edgee.send(
-    model="gpt-4o",
+    model="mistral/mistral-small-latest",
     input={
         "messages": [
             {"role": "system", "content": "You are a helpful assistant."},
@@ -35,29 +35,60 @@
 print()
 
 # Test 3: With tools
-print("Test 3: With tools")
-response3 = edgee.send(
-    model="gpt-4o",
+#print("Test 3: With tools")
+#response3 = edgee.send(
+#    model="gpt-4o",
+#    input={
+#        "messages": [{"role": "user", "content": "What is the weather in Paris?"}],
+#        "tools": [
+#            {
+#                "type": "function",
+#                "function": {
+#                    "name": "get_weather",
+#                    "description": "Get the current weather for a location",
+#                    "parameters": {
+#                        "type": "object",
+#                        "properties": {
+#                            "location": {"type": "string", "description": "City name"},
+#                        },
+#                        "required": ["location"],
+#                    },
+#                },
+#            },
+#        ],
+#        "tool_choice": "auto",
+#    },
+#)
+#print(f"Content: {response3.choices[0].message.get('content')}")
+#print(f"Tool calls: {response3.choices[0].message.get('tool_calls')}")
+#print()
+
+# Test 4: Streaming
+print("Test 4: Streaming")
+for chunk in edgee.stream(
+    model="mistral/mistral-small-latest",
+    input="Tell me a short story about a robot",
+):
+    if chunk.choices[0].delta.content:
+        print(chunk.choices[0].delta.content, end="", flush=True)
+print("\n")
+
+# Test 5: Streaming with messages
+print("Test 5: Streaming with system message")
+for chunk in edgee.stream(
+    model="mistral/mistral-small-latest",
     input={
-        "messages": [{"role": "user", "content": "What is the weather in Paris?"}],
-        "tools": [
-            {
-                "type": "function",
-                "function": {
-                    "name": "get_weather",
-                    "description": "Get the current weather for a location",
-                    "parameters": {
-                        "type": "object",
-                        "properties": {
-                            "location": {"type": "string", "description": "City name"},
-                        },
-                        "required": ["location"],
-                    },
-                },
-            },
+        "messages": [
+            {"role": "system", "content": "You are a poetic assistant. Respond in rhyme."},
+            {"role": "user", "content": "What is Python?"},
         ],
-        "tool_choice": "auto",
     },
-)
-print(f"Content: {response3.choices[0].message.get('content')}")
-print(f"Tool calls: {response3.choices[0].message.get('tool_calls')}")
+):
+    delta = chunk.choices[0].delta
+    if delta.role:
+        print(f"[Role: {delta.role}]")
+    if delta.content:
+        print(delta.content, end="", flush=True)
+    if chunk.choices[0].finish_reason:
+        print(f"\n[Finished: {chunk.choices[0].finish_reason}]")
+print()