Skip to content

Commit 796cf6b

Browse files
CLEMENTINATORCLEMENTINATOR
committed
add helper properties
1 parent 3a784f3 commit 796cf6b

3 files changed

Lines changed: 107 additions & 84 deletions

File tree

‎README.md‎

Lines changed: 33 additions & 30 deletions
Original file line numberDiff line numberDiff line change
@@ -73,47 +73,45 @@ if response.choices[0].message.get("tool_calls"):
7373

7474
### Streaming
7575

76-
You can enable streaming by setting `stream=True` in the `send()` method, or by using the convenience `stream()` method.
76+
#### Simple Text Streaming
7777

78-
#### Using send(stream=True)
78+
The simplest way to stream text responses:
7979

8080
```python
81-
for chunk in edgee.send(
82-
model="gpt-4o",
83-
input="Tell me a story",
84-
stream=True,
85-
):
86-
if chunk.choices[0].delta.content:
87-
print(chunk.choices[0].delta.content, end="", flush=True)
81+
for text in edgee.stream_text(model="gpt-4o", input="Tell me a story"):
82+
print(text, end="", flush=True)
8883
```
8984

90-
#### Using stream() method
85+
#### Streaming with More Control
86+
87+
Access chunk properties when you need more control:
9188

9289
```python
93-
for chunk in edgee.stream(
94-
model="gpt-4o",
95-
input="Tell me a story",
96-
):
97-
if chunk.choices[0].delta.content:
98-
print(chunk.choices[0].delta.content, end="", flush=True)
90+
for chunk in edgee.stream(model="gpt-4o", input="Tell me a story"):
91+
if chunk.text:
92+
print(chunk.text, end="", flush=True)
9993
```
10094

101-
### Streaming with Messages
95+
#### Alternative: Using send(stream=True)
10296

10397
```python
104-
for chunk in edgee.send(
105-
model="gpt-4o",
106-
input={
107-
"messages": [
108-
{"role": "system", "content": "You are a helpful assistant."},
109-
{"role": "user", "content": "Hello!"},
110-
],
111-
},
112-
stream=True,
113-
):
114-
delta = chunk.choices[0].delta
115-
if delta.content:
116-
print(delta.content, end="", flush=True)
98+
for chunk in edgee.send(model="gpt-4o", input="Tell me a story", stream=True):
99+
if chunk.text:
100+
print(chunk.text, end="", flush=True)
101+
```
102+
103+
#### Accessing Full Chunk Data
104+
105+
When you need complete access to the streaming response:
106+
107+
```python
108+
for chunk in edgee.stream(model="gpt-4o", input="Hello"):
109+
if chunk.role:
110+
print(f"Role: {chunk.role}")
111+
if chunk.text:
112+
print(chunk.text, end="", flush=True)
113+
if chunk.finish_reason:
114+
print(f"\nFinish: {chunk.finish_reason}")
117115
```
118116

119117
## Response
@@ -144,6 +142,11 @@ class Usage:
144142
class StreamChunk:
145143
choices: list[StreamChoice]
146144

145+
# Convenience properties for easy access
146+
text: str | None # Shortcut for choices[0].delta.content
147+
role: str | None # Shortcut for choices[0].delta.role
148+
finish_reason: str | None # Shortcut for choices[0].finish_reason
149+
147150
@dataclass
148151
class StreamChoice:
149152
index: int

‎edgee/__init__.py‎

Lines changed: 38 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -85,6 +85,27 @@ class StreamChoice:
8585
class StreamChunk:
8686
choices: list[StreamChoice]
8787

88+
@property
89+
def text(self) -> str | None:
90+
"""Convenience property to get text content from the first choice."""
91+
if self.choices and self.choices[0].delta.content:
92+
return self.choices[0].delta.content
93+
return None
94+
95+
@property
96+
def role(self) -> str | None:
97+
"""Convenience property to get role from the first choice."""
98+
if self.choices and self.choices[0].delta.role:
99+
return self.choices[0].delta.role
100+
return None
101+
102+
@property
103+
def finish_reason(self) -> str | None:
104+
"""Convenience property to get finish_reason from the first choice."""
105+
if self.choices and self.choices[0].finish_reason:
106+
return self.choices[0].finish_reason
107+
return None
108+
88109

89110
@dataclass
90111
class EdgeeConfig:
@@ -256,3 +277,20 @@ def stream(
256277
Yields StreamChunk objects as they arrive from the API.
257278
"""
258279
return self.send(model=model, input=input, stream=True)
280+
281+
def stream_text(
282+
self,
283+
model: str,
284+
input: str | InputObject | dict,
285+
):
286+
"""Stream only the text content from the completion.
287+
288+
Convenience method that yields only non-empty text strings,
289+
filtering out chunks without content.
290+
291+
Yields:
292+
str: Text content from each chunk
293+
"""
294+
for chunk in self.stream(model=model, input=input):
295+
if chunk.text:
296+
yield chunk.text

‎example/test.py‎

Lines changed: 36 additions & 54 deletions
Original file line numberDiff line numberDiff line change
@@ -35,61 +35,43 @@
3535
print()
3636

3737
# Test 3: With tools
38-
#print("Test 3: With tools")
39-
#response3 = edgee.send(
40-
# model="gpt-4o",
41-
# input={
42-
# "messages": [{"role": "user", "content": "What is the weather in Paris?"}],
43-
# "tools": [
44-
# {
45-
# "type": "function",
46-
# "function": {
47-
# "name": "get_weather",
48-
# "description": "Get the current weather for a location",
49-
# "parameters": {
50-
# "type": "object",
51-
# "properties": {
52-
# "location": {"type": "string", "description": "City name"},
53-
# },
54-
# "required": ["location"],
55-
# },
56-
# },
57-
# },
58-
# ],
59-
# "tool_choice": "auto",
60-
# },
61-
#)
62-
#print(f"Content: {response3.choices[0].message.get('content')}")
63-
#print(f"Tool calls: {response3.choices[0].message.get('tool_calls')}")
64-
#print()
65-
66-
# Test 4: Streaming with send(stream=True)
67-
print("Test 4: Streaming with send(stream=True)")
68-
for chunk in edgee.send(
69-
model="mistral/mistral-small-latest",
70-
input="Tell me a short story about a robot",
71-
stream=True,
72-
):
73-
if chunk.choices[0].delta.content:
74-
print(chunk.choices[0].delta.content, end="", flush=True)
75-
print("\n")
76-
77-
# Test 5: Streaming with messages
78-
print("Test 5: Streaming with system message")
79-
for chunk in edgee.stream(
80-
model="mistral/mistral-small-latest",
38+
print("Test 3: With tools")
39+
response3 = edgee.send(
40+
model="gpt-4o",
8141
input={
82-
"messages": [
83-
{"role": "system", "content": "You are a poetic assistant. Respond in rhyme."},
84-
{"role": "user", "content": "What is Python?"},
42+
"messages": [{"role": "user", "content": "What is the weather in Paris?"}],
43+
"tools": [
44+
{
45+
"type": "function",
46+
"function": {
47+
"name": "get_weather",
48+
"description": "Get the current weather for a location",
49+
"parameters": {
50+
"type": "object",
51+
"properties": {
52+
"location": {"type": "string", "description": "City name"},
53+
},
54+
"required": ["location"],
55+
},
56+
},
57+
},
8558
],
59+
"tool_choice": "auto",
8660
},
87-
):
88-
delta = chunk.choices[0].delta
89-
if delta.role:
90-
print(f"[Role: {delta.role}]")
91-
if delta.content:
92-
print(delta.content, end="", flush=True)
93-
if chunk.choices[0].finish_reason:
94-
print(f"\n[Finished: {chunk.choices[0].finish_reason}]")
61+
)
62+
print(f"Content: {response3.choices[0].message.get('content')}")
63+
print(f"Tool calls: {response3.choices[0].message.get('tool_calls')}")
9564
print()
65+
66+
# Test 4: Streaming (simplest way)
67+
print("Test 4: Streaming (simplest way)")
68+
for text in edgee.stream_text(model="mistral/mistral-small-latest", input="Tell me a short story about a robot"):
69+
print(text, end="", flush=True)
70+
print("\n")
71+
72+
# Test 5: Streaming with more control
73+
print("Test 5: Streaming with more control")
74+
for chunk in edgee.stream(model="mistral/mistral-small-latest", input="What is Python?"):
75+
if chunk.text:
76+
print(chunk.text, end="", flush=True)
77+
print("\n")

0 commit comments

Comments
 (0)