Spaces:
Running
Running
| 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 | |