Skip to content

Commit 2325fba

Browse files
UtkarshTheDevUtkarshTheDev
committed
Improved steaming errors handling in client packages
1 parent 58888f3 commit 2325fba

File tree

3 files changed

+75
-71
lines changed

3 files changed

+75
-71
lines changed

‎client/python_client/examples/test_client.py

Lines changed: 10 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -16,11 +16,17 @@ async def main():
1616
print(response.text)
1717
print()
1818

19-
# Test streaming with proper spacing
19+
# Test streaming with proper spacing and error handling
2020
print("Streaming response:")
21-
async for token in client.stream_generate("What are you working on? Tell me about your current project."):
22-
print(token, end="", flush=True)
23-
print("\n")
21+
try:
22+
async for token in client.stream_generate("What are you working on? Tell me about your current project."):
23+
if token.startswith("Error:"):
24+
print(f"\n{token}")
25+
break
26+
print(token, end="", flush=True)
27+
print("\n")
28+
except Exception as e:
29+
print(f"\nStreaming error: {str(e)}")
2430

2531
# Test chat with context
2632
messages = [

‎client/python_client/locallab/__pycache__/client.cpython-313.pyc

-99 Bytes
Binary file not shown.

‎client/python_client/locallab/client.py

Lines changed: 65 additions & 67 deletions
Original file line numberDiff line numberDiff line change
@@ -186,81 +186,79 @@ async def stream_generate(self, prompt: str, options: Optional[Union[GenerateOpt
186186

187187
async with self.session.post("/generate", json=data) as response:
188188
if response.status != 200:
189-
try:
190-
error_data = await response.json()
191-
error_msg = error_data.get("detail", "Streaming failed")
192-
logger.error(f"Streaming error: {error_msg}")
193-
yield f"\nError: {error_msg}"
194-
return
195-
except:
196-
yield "\nError: Streaming failed"
197-
return
189+
error_msg = await response.text()
190+
logger.error(f"Streaming error: {error_msg}")
191+
yield f"Error: {error_msg}"
192+
return
198193

199194
buffer = ""
200195
current_sentence = ""
201196
last_token_was_space = False
202197

203-
async for line in response.content:
204-
if line:
205-
try:
206-
line = line.decode('utf-8').strip()
207-
# Skip empty lines
208-
if not line:
209-
continue
210-
211-
# Handle SSE format
212-
if line.startswith("data: "):
213-
line = line[6:] # Remove "data: " prefix
214-
215-
# Skip control messages
216-
if line in ["[DONE]", "[ERROR]"]:
217-
continue
218-
198+
try:
199+
async for line in response.content:
200+
if line:
219201
try:
220-
# Try to parse as JSON
221-
data = json.loads(line)
222-
text = data.get("text", data.get("response", ""))
223-
except json.JSONDecodeError:
224-
# If not JSON, use the line as is
225-
text = line
226-
227-
if text:
228-
# Clean up any special tokens
229-
text = text.replace("<|", "").replace("|>", "")
230-
text = text.replace("<", "").replace(">", "")
231-
text = text.replace("[", "").replace("]", "")
232-
text = text.replace("{", "").replace("}", "")
233-
text = text.replace("data:", "")
234-
text = text.replace("��", "")
235-
text = text.replace("\\n", "\n")
236-
text = text.replace("|user|", "")
237-
text = text.replace("|The", "The")
238-
text = text.replace("/|assistant|", "").replace("/|user|", "")
239-
240-
# Add space between words if needed
241-
if (not text.startswith(" ") and
242-
not text.startswith("\n") and
243-
not last_token_was_space and
244-
buffer and
245-
not buffer.endswith(" ") and
246-
not buffer.endswith("\n")):
247-
text = " " + text
202+
line = line.decode('utf-8').strip()
203+
# Skip empty lines
204+
if not line:
205+
continue
248206

249-
# Update tracking variables
250-
buffer += text
251-
current_sentence += text
252-
last_token_was_space = text.endswith(" ") or text.endswith("\n")
253-
254-
# Check for sentence completion
255-
if any(current_sentence.endswith(p) for p in [".", "!", "?", "\n"]):
256-
current_sentence = ""
257-
258-
yield text
207+
# Handle SSE format
208+
if line.startswith("data: "):
209+
line = line[6:] # Remove "data: " prefix
210+
211+
# Skip control messages
212+
if line in ["[DONE]", "[ERROR]"]:
213+
continue
214+
215+
try:
216+
# Try to parse as JSON
217+
data = json.loads(line)
218+
text = data.get("text", data.get("response", ""))
219+
except json.JSONDecodeError:
220+
# If not JSON, use the line as is
221+
text = line
222+
223+
if text:
224+
# Clean up any special tokens
225+
text = text.replace("<|", "").replace("|>", "")
226+
text = text.replace("<", "").replace(">", "")
227+
text = text.replace("[", "").replace("]", "")
228+
text = text.replace("{", "").replace("}", "")
229+
text = text.replace("data:", "")
230+
text = text.replace("��", "")
231+
text = text.replace("\\n", "\n")
232+
text = text.replace("|user|", "")
233+
text = text.replace("|The", "The")
234+
text = text.replace("/|assistant|", "").replace("/|user|", "")
235+
text = text.replace("assistant", "").replace("Error:", "")
236+
237+
# Add space between words if needed
238+
if (not text.startswith(" ") and
239+
not text.startswith("\n") and
240+
not last_token_was_space and
241+
buffer and
242+
not buffer.endswith(" ") and
243+
not buffer.endswith("\n")):
244+
text = " " + text
245+
246+
# Update tracking variables
247+
buffer += text
248+
current_sentence += text
249+
last_token_was_space = text.endswith(" ") or text.endswith("\n")
250+
251+
yield text
252+
253+
except Exception as e:
254+
logger.error(f"Error processing stream chunk: {str(e)}")
255+
yield f"\nError: {str(e)}"
256+
return
259257

260-
except Exception as e:
261-
logger.error(f"Error processing stream chunk: {str(e)}")
262-
yield f"\nError: {str(e)}"
263-
return
258+
except Exception as e:
259+
logger.error(f"Stream connection error: {str(e)}")
260+
yield f"\nError: Connection error - {str(e)}"
261+
return
264262

265263
async def generate(self, prompt: str, options: Optional[Union[GenerateOptions, Dict]] = None) -> GenerateResponse:
266264
"""Generate text from prompt"""

0 commit comments

Comments
 (0)