import gradio as gr import httpx import json import asyncio import os from chat_state import chat_state from config import OLLAMA_URL, DEFAULT_TEMPERATURE, DEFAULT_SYSTEM_MESSAGE theme = gr.themes.Soft( primary_hue="yellow", neutral_hue="neutral", text_size="md", spacing_size="md", radius_size="md", font=[gr.themes.GoogleFont('Montserrat'), gr.themes.GoogleFont('ui-sans-serif'), 'system-ui', 'sans-serif'], ) async def fetch_available_models(): async with httpx.AsyncClient() as client: try: response = await client.get(f"{OLLAMA_URL}/api/tags") response.raise_for_status() data = response.json() return [model["name"] for model in data.get("models", [])] except httpx.HTTPStatusError as e: print(f"Error fetching models: {e}") return [] async def get_model_info(model_name): async with httpx.AsyncClient() as client: try: response = await client.post(f"{OLLAMA_URL}/api/show", json={"name": model_name}) response.raise_for_status() return response.json() except httpx.HTTPStatusError as e: print(f"Error fetching model info: {e}") return {} async def call_ollama_api(prompt, history): messages = [{"role": "system", "content": chat_state.system_message}] for user_msg, assistant_msg in history: if user_msg: messages.append({"role": "user", "content": user_msg}) if assistant_msg: messages.append({"role": "assistant", "content": assistant_msg}) messages.append({"role": "user", "content": prompt}) payload = { "model": chat_state.model, "messages": messages, "stream": True, "temperature": chat_state.temperature } async with httpx.AsyncClient() as client: try: async with client.stream("POST", f"{OLLAMA_URL}/api/chat", json=payload, timeout=30.0) as response: response.raise_for_status() full_response = "" async for line in response.aiter_lines(): if line: json_line = json.loads(line) message_content = json_line.get('message', {}).get('content', '') if message_content: full_response += message_content yield full_response if json_line.get('done'): break except httpx.HTTPStatusError as e: yield f"An error occurred: {e}" except asyncio.TimeoutError: yield "The request timed out. Please try again." async def user(user_message, history): return "", history + [[user_message, None]] async def bot(history): user_message = history[-1][0] bot_message_generator = call_ollama_api(user_message, history[:-1]) async for message_content in bot_message_generator: history[-1][1] = message_content yield history def clear_chat(): return None def save_chat_history(history, filename="chat_history.json"): with open(filename, "w") as f: json.dump(history, f) return f"Chat history saved to {filename}" def load_chat_history(filename="chat_history.json"): try: with open(filename, "r") as f: return json.load(f) except FileNotFoundError: return None async def change_model(model_name): chat_state.model = model_name model_info = await get_model_info(model_name) info_text = f"Model: {model_name}\n" info_text += f"Parameter Size: {model_info.get('details', {}).get('parameter_size', 'Unknown')}\n" info_text += f"Quantization: {model_info.get('details', {}).get('quantization_level', 'Unknown')}\n" info_text += f"Format: {model_info.get('details', {}).get('format', 'Unknown')}" return f"Model changed to {chat_state.model}", info_text def update_temperature(new_temp): chat_state.temperature = float(new_temp) return f"Temperature set to {chat_state.temperature}" def update_system_message(new_message): chat_state.system_message = new_message return f"System message updated: {chat_state.system_message}" async def initialize_interface(): chat_state.available_models = await fetch_available_models() with gr.Blocks(theme=theme) as demo: gr.Markdown("# 🤖 Enhanced Ollama Chatbot Interface") with gr.Row(): with gr.Column(scale=7): chatbot = gr.Chatbot(height=600, elem_id="chatbot") with gr.Row(): msg = gr.Textbox( label="Message", placeholder="Type your message here...", scale=4, elem_id="user-input" ) send = gr.Button("Send", scale=1, elem_id="send-btn") with gr.Column(scale=3): with gr.Accordion("Model Settings", open=True): model_dropdown = gr.Dropdown( choices=chat_state.available_models, label="Select Model", value=chat_state.available_models[0] if chat_state.available_models else None, elem_id="model-dropdown" ) model_info = gr.Textbox(label="Model Information", interactive=False, lines=4) temp_slider = gr.Slider( minimum=0, maximum=1, value=DEFAULT_TEMPERATURE, step=0.1, label="Temperature", elem_id="temp-slider" ) with gr.Accordion("System Message", open=False): system_message_input = gr.Textbox( label="System Message", value=DEFAULT_SYSTEM_MESSAGE, lines=3, elem_id="system-message" ) update_system_button = gr.Button("Update System Message", elem_id="update-system-btn") with gr.Accordion("Chat Management", open=False): with gr.Row(): clear = gr.Button("Clear Chat", elem_id="clear-btn") save_button = gr.Button("Save Chat", elem_id="save-btn") load_button = gr.Button("Load Chat", elem_id="load-btn") status_box = gr.Textbox(label="Status", interactive=False, elem_id="status-box") # Event handlers send_event = msg.submit(user, [msg, chatbot], [msg, chatbot], queue=False).then( bot, chatbot, chatbot ) send.click(user, [msg, chatbot], [msg, chatbot], queue=False).then( bot, chatbot, chatbot ) clear.click(clear_chat, outputs=[chatbot]) model_dropdown.change(change_model, inputs=[model_dropdown], outputs=[status_box, model_info]) temp_slider.change(update_temperature, inputs=[temp_slider], outputs=[status_box]) update_system_button.click(update_system_message, inputs=[system_message_input], outputs=[status_box]) save_button.click(save_chat_history, inputs=[chatbot], outputs=[status_box]) load_button.click(load_chat_history, outputs=[chatbot]) # Initialize the first model if chat_state.available_models: chat_state.model = chat_state.available_models[0] return demo if __name__ == "__main__": demo = asyncio.run(initialize_interface()) demo.launch(share=True)