Skip to content

Commit 66f5739

Browse files
committed
add helper properties
1 parent 3a784f3 commit 66f5739

3 files changed

Lines changed: 148 additions & 89 deletions

File tree

README.md

Lines changed: 42 additions & 33 deletions
Original file line numberDiff line numberDiff line change
@@ -24,7 +24,7 @@ response = edgee.send(
2424
input="What is the capital of France?",
2525
)
2626

27-
print(response.choices[0].message["content"])
27+
print(response.text)
2828
```
2929

3030
### Full Input with Messages
@@ -67,53 +67,51 @@ response = edgee.send(
6767
},
6868
)
6969

70-
if response.choices[0].message.get("tool_calls"):
71-
print(response.choices[0].message["tool_calls"])
70+
if response.tool_calls:
71+
print(response.tool_calls)
7272
```
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
@@ -124,6 +122,12 @@ class SendResponse:
124122
choices: list[Choice]
125123
usage: Optional[Usage]
126124

125+
# Convenience properties for easy access
126+
text: str | None # Shortcut for choices[0].message["content"]
127+
message: dict | None # Shortcut for choices[0].message
128+
finish_reason: str | None # Shortcut for choices[0].finish_reason
129+
tool_calls: list | None # Shortcut for choices[0].message["tool_calls"]
130+
127131
@dataclass
128132
class Choice:
129133
index: int
@@ -144,6 +148,11 @@ class Usage:
144148
class StreamChunk:
145149
choices: list[StreamChoice]
146150

151+
# Convenience properties for easy access
152+
text: str | None # Shortcut for choices[0].delta.content
153+
role: str | None # Shortcut for choices[0].delta.role
154+
finish_reason: str | None # Shortcut for choices[0].finish_reason
155+
147156
@dataclass
148157
class StreamChoice:
149158
index: int

edgee/__init__.py

Lines changed: 66 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -66,6 +66,34 @@ class SendResponse:
6666
choices: list[Choice]
6767
usage: Usage | None = None
6868

69+
@property
70+
def text(self) -> str | None:
71+
"""Convenience property to get text content from the first choice."""
72+
if self.choices and self.choices[0].message.get("content"):
73+
return self.choices[0].message["content"]
74+
return None
75+
76+
@property
77+
def message(self) -> dict | None:
78+
"""Convenience property to get the message from the first choice."""
79+
if self.choices:
80+
return self.choices[0].message
81+
return None
82+
83+
@property
84+
def finish_reason(self) -> str | None:
85+
"""Convenience property to get finish_reason from the first choice."""
86+
if self.choices and self.choices[0].finish_reason:
87+
return self.choices[0].finish_reason
88+
return None
89+
90+
@property
91+
def tool_calls(self) -> list | None:
92+
"""Convenience property to get tool_calls from the first choice."""
93+
if self.choices and self.choices[0].message.get("tool_calls"):
94+
return self.choices[0].message["tool_calls"]
95+
return None
96+
6997

7098
@dataclass
7199
class StreamDelta:
@@ -85,6 +113,27 @@ class StreamChoice:
85113
class StreamChunk:
86114
choices: list[StreamChoice]
87115

116+
@property
117+
def text(self) -> str | None:
118+
"""Convenience property to get text content from the first choice."""
119+
if self.choices and self.choices[0].delta.content:
120+
return self.choices[0].delta.content
121+
return None
122+
123+
@property
124+
def role(self) -> str | None:
125+
"""Convenience property to get role from the first choice."""
126+
if self.choices and self.choices[0].delta.role:
127+
return self.choices[0].delta.role
128+
return None
129+
130+
@property
131+
def finish_reason(self) -> str | None:
132+
"""Convenience property to get finish_reason from the first choice."""
133+
if self.choices and self.choices[0].finish_reason:
134+
return self.choices[0].finish_reason
135+
return None
136+
88137

89138
@dataclass
90139
class EdgeeConfig:
@@ -256,3 +305,20 @@ def stream(
256305
Yields StreamChunk objects as they arrive from the API.
257306
"""
258307
return self.send(model=model, input=input, stream=True)
308+
309+
def stream_text(
310+
self,
311+
model: str,
312+
input: str | InputObject | dict,
313+
):
314+
"""Stream only the text content from the completion.
315+
316+
Convenience method that yields only non-empty text strings,
317+
filtering out chunks without content.
318+
319+
Yields:
320+
str: Text content from each chunk
321+
"""
322+
for chunk in self.stream(model=model, input=input):
323+
if chunk.text:
324+
yield chunk.text

example/test.py

Lines changed: 40 additions & 56 deletions
Original file line numberDiff line numberDiff line change
@@ -16,7 +16,7 @@
1616
model="mistral/mistral-small-latest",
1717
input="What is the capital of France?",
1818
)
19-
print(f"Content: {response1.choices[0].message['content']}")
19+
print(f"Content: {response1.text}")
2020
print(f"Usage: {response1.usage}")
2121
print()
2222

@@ -31,65 +31,49 @@
3131
],
3232
},
3333
)
34-
print(f"Content: {response2.choices[0].message['content']}")
34+
print(f"Content: {response2.text}")
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.text}")
63+
print(f"Tool calls: {response3.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(
69+
model="mistral/mistral-small-latest", input="Tell me a short story about a robot"
70+
):
71+
print(text, end="", flush=True)
72+
print("\n")
73+
74+
# Test 5: Streaming with more control
75+
print("Test 5: Streaming with more control")
76+
for chunk in edgee.stream(model="mistral/mistral-small-latest", input="What is Python?"):
77+
if chunk.text:
78+
print(chunk.text, end="", flush=True)
79+
print("\n")

0 commit comments

Comments
 (0)