Spaces:
Running
Running
File size: 9,746 Bytes
2a2c2ad fcb8f25 2a2c2ad fcb8f25 2a2c2ad fcb8f25 2a2c2ad fcb8f25 2a2c2ad fcb8f25 2a2c2ad fcb8f25 2a2c2ad fcb8f25 2a2c2ad fcb8f25 2a2c2ad fcb8f25 2a2c2ad fcb8f25 2a2c2ad fcb8f25 2a2c2ad fcb8f25 2a2c2ad fcb8f25 2a2c2ad fcb8f25 2a2c2ad fcb8f25 2a2c2ad fcb8f25 2a2c2ad fcb8f25 2a2c2ad fcb8f25 2a2c2ad fcb8f25 2a2c2ad fcb8f25 2a2c2ad fcb8f25 2a2c2ad fcb8f25 2a2c2ad fcb8f25 2a2c2ad fcb8f25 2a2c2ad fcb8f25 2a2c2ad fcb8f25 2a2c2ad fcb8f25 2a2c2ad fcb8f25 2a2c2ad fcb8f25 2a2c2ad fcb8f25 2a2c2ad fcb8f25 2a2c2ad fcb8f25 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 |
import json
import requests
from rich.console import Console
from rich.live import Live
from rich.markdown import Markdown
from rich.panel import Panel
from rich.text import Text
class StreamResponseHandler:
"""
A utility class for handling streaming responses from API endpoints.
Provides rich formatting and real-time updates of the response content.
"""
def __init__(self, console=None):
"""
Initialize the stream response handler.
Args:
console (Console, optional): A Rich console instance. If not provided, a new one will be created.
"""
self.console = console or Console()
def check_server_health(self, health_url="http://localhost:8000/health"):
"""
Check if the server is running and accessible.
Args:
health_url (str, optional): The URL to check server health. Defaults to "http://localhost:8000/health".
Returns:
bool: True if the server is running and accessible, False otherwise.
"""
try:
self.console.print("Checking server health...", style="bold yellow")
response = requests.get(health_url)
if response.status_code == 200:
self.console.print("[bold green]✓ Server is running and accessible.[/]")
return True
else:
self.console.print(
f"[bold red]✗ Server health check failed[/] with status code: {response.status_code}"
)
return False
except requests.exceptions.ConnectionError:
self.console.print(
"[bold red]✗ Error:[/] Could not connect to the server. Make sure it's running."
)
return False
except Exception as e:
self.console.print(f"[bold red]✗ Error checking server health:[/] {e}")
return False
def stream_response(
self, url, payload=None, params=None, method="POST", title="AI Response"
):
"""
Send a request to an endpoint and stream the output to the terminal.
Args:
url (str): The URL of the endpoint to send the request to.
payload (dict, optional): The JSON payload to send in the request body. Defaults to None.
params (dict, optional): The query parameters to send in the request. Defaults to None.
method (str, optional): The HTTP method to use. Defaults to "POST".
title (str, optional): The title to display in the panel. Defaults to "AI Response".
Returns:
bool: True if the streaming was successful, False otherwise.
"""
# Display request information
self.console.print(f"Sending request to [bold cyan]{url}[/]")
if payload:
self.console.print("Payload:", style="bold")
self.console.print(json.dumps(payload, indent=2))
if params:
self.console.print("Parameters:", style="bold")
self.console.print(json.dumps(params, indent=2))
try:
# Prepare the request
request_kwargs = {"stream": True}
if payload:
request_kwargs["json"] = payload
if params:
request_kwargs["params"] = params
# Make the request
with getattr(requests, method.lower())(url, **request_kwargs) as response:
# Check if the request was successful
if response.status_code != 200:
self.console.print(
f"[bold red]Error:[/] Received status code {response.status_code}"
)
self.console.print(f"Response: {response.text}")
return False
# Initialize an empty response text
full_response = ""
# Use Rich's Live display to update the content in place
with Live(
Panel("Waiting for response...", title=title, border_style="blue"),
refresh_per_second=10,
) as live:
# Process the streaming response
for line in response.iter_lines():
if line:
# Decode the line and parse it as JSON
decoded_line = line.decode("utf-8")
try:
# Parse the JSON
data = json.loads(decoded_line)
# Extract and display the content
if isinstance(data, dict):
if "content" in data:
for content in data["content"]:
if content.get("type") == "text":
text_content = content.get("text", "")
# Append to the full response
full_response += text_content
# Update the live display with the current full response
live.update(
Panel(
Markdown(full_response),
title=title,
border_style="green",
)
)
elif content.get("type") == "image_url":
image_url = content.get(
"image_url", {}
).get("url", "")
# Add a note about the image URL
image_note = (
f"\n\n[Image URL: {image_url}]"
)
full_response += image_note
live.update(
Panel(
Markdown(full_response),
title=title,
border_style="green",
)
)
elif "edited_image_url" in data:
# Handle edited image URL from edit endpoint
image_url = data.get("edited_image_url", "")
image_note = (
f"\n\n[Edited Image URL: {image_url}]"
)
full_response += image_note
live.update(
Panel(
Markdown(full_response),
title=title,
border_style="green",
)
)
else:
# For other types of data, just show the JSON
live.update(
Panel(
Text(json.dumps(data, indent=2)),
title="Raw JSON Response",
border_style="yellow",
)
)
else:
live.update(
Panel(
Text(decoded_line),
title="Raw Response",
border_style="yellow",
)
)
except json.JSONDecodeError:
# If it's not valid JSON, just show the raw line
live.update(
Panel(
Text(f"Raw response: {decoded_line}"),
title="Invalid JSON",
border_style="red",
)
)
self.console.print("[bold green]Stream completed.[/]")
return True
except requests.exceptions.ConnectionError:
self.console.print(
f"[bold red]Error:[/] Could not connect to the server at {url}",
style="red",
)
self.console.print(
"Make sure the server is running and accessible.", style="red"
)
return False
except requests.exceptions.RequestException as e:
self.console.print(f"[bold red]Error:[/] {e}", style="red")
return False
|