Spaces:
Running
Running
File size: 7,682 Bytes
2a2c2ad |
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 |
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 |