File size: 7,190 Bytes
fe52ef9 | 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 | import gradio as gr
import json
import os
from typing import List, Dict, Any, Union, Generator
from generators import chat_completion_stream
from client_utils import engine_map
from file_processors import FileProcessorFactory
import logging
import time
# Configure logging first
logging.basicConfig(
level=logging.INFO,
format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
)
# Then get a logger
logger = logging.getLogger(__name__)
class Chatbot:
def __init__(self):
"""Initialize a new chatbot with history management."""
self._stop = False # Flag to interrupt streaming
self.engine_map = engine_map
self.content_exclusions = ["<|eos|>"]
def process_file(self, files: Union[List[str], str]) -> Union[str, List[Dict[str, Any]]]:
"""Process uploaded files and return their contents as JSON."""
if isinstance(files, str):
files = [files]
file_contents = []
for file_path in files:
try:
processor = FileProcessorFactory.get_processor(file_path)
file_contents.append(processor.process(file_path))
except Exception as e:
file_contents.append(f"Error processing file {os.path.basename(file_path)}: {str(e)}\n")
return json.dumps(file_contents)
def process_message(self, message: Union[str, dict]) -> dict:
"""Convert various message formats into a standardized message dictionary."""
if isinstance(message, dict):
if message.get("text") and message.get("files"):
return {
"role": "user",
"content": f'{message["text"]}\n\nFile Upload(s):\n{self.process_file(message["files"])}'
}
elif message.get("text"):
return {
"role": "user",
"content": message["text"]
}
elif message.get("files"):
return {
"role": "user",
"content": f'File Upload(s):\n{self.process_file(message["files"])}'
}
elif message.get("content") and message.get("role"):
if isinstance(message["content"], tuple):
return {
"role": message["role"],
"content": f'File Upload(s):\n{self.process_file(message["content"][0])}'
}
elif isinstance(message["content"], str):
return {
"role": message["role"],
"content": message["content"]
}
elif isinstance(message, str):
if message:
return {
"role": "user",
"content": message,
}
return
def predict(self, message: Union[str, dict], history: List[Dict[str, Any]], engine: str, temperature: float) -> Generator:
"""Generate a streaming response for the given message and conversation history."""
self._stop = False # Reset stop flag
yield "", gr.update(interactive=False, submit_btn=False, stop_btn=True)
history.append(message)
# Process all messages and update the list
history = [msg for msg in (self.process_message(msg) for msg in history) if msg is not None]
response = []
streamed_content = "" # Accumulate the assistant's response here
def fade_html(text: str, fade_count: int = 200) -> str:
"""
Fade all characters if text is shorter than fade_count.
Otherwise, fade only the last fade_count characters.
"""
if not text:
return text
text_len = len(text)
if text_len <= fade_count:
# Fade all characters
fade_spans = ''
for i, c in enumerate(text):
opacity = 1.0 - (i / max(text_len - 1, 1))
fade_spans += f'<span style="opacity:{opacity:.2f};transition:opacity 0.2s">{c}</span>'
return fade_spans
else:
base = text[:-fade_count]
fades = text[-fade_count:]
fade_spans = ''
for i, c in enumerate(fades):
opacity = 1.0 - (i / (fade_count - 1))
fade_spans += f'<span style="opacity:{opacity:.2f};transition:opacity 0.2s">{c}</span>'
return base + fade_spans
for chunk in chat_completion_stream(history, engine=engine, temperature=temperature):
logger.debug("Chunk received: %s", chunk)
if self._stop:
break
try:
content = chunk.get("content", "")
if chunk.get("metadata"):
# Do not stream metadata contents, just update/append as before
new_metadata = True
for i, resp in enumerate(response):
if resp.get("metadata", {}).get("id") == chunk["metadata"]["id"]:
response[i]["content"] = content
response[i]["metadata"] = chunk["metadata"]
new_metadata = False
break
if new_metadata:
response.append({
"role": "assistant",
"content": content,
"metadata": chunk["metadata"]
})
# Do not stream char-by-char for metadata
# yield response, gr.skip()
else:
# Stream each character in the latest assistant message (no metadata)
if not response or response[-1].get("metadata"):
response.append({
"role": "assistant",
"content": content,
})
streamed_content = "" # Reset for new message
else:
response[-1]["content"] = content
latest = response[-1]
prev_len = len(streamed_content)
new_content = latest.get("content", "")
for c in new_content[prev_len:]:
if self._stop:
break
streamed_content += c
# Only fade the last fade_count characters of the current message
latest["content"] = fade_html(streamed_content)
yield response, gr.skip()
time.sleep(0.003)
except Exception as e:
logger.info(f"Error processing chunk from chat completion: {e}")
continue
yield response, gr.update(interactive=True, submit_btn=True, stop_btn=False)
def stop(self):
"""Interrupt the currently streaming response."""
self._stop = True |