File size: 3,868 Bytes
204d759
 
b293ab8
77ed406
94ef460
 
204d759
bc2d859
77ed406
 
5a2e9d4
 
 
 
 
bc2d859
94ef460
77ed406
bc2d859
204d759
 
 
bc2d859
204d759
 
 
 
 
 
 
 
 
bc2d859
204d759
 
 
bc2d859
204d759
 
bc2d859
204d759
 
5a2e9d4
204d759
 
5a2e9d4
204d759
 
 
 
 
 
 
 
5a2e9d4
204d759
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
5a2e9d4
 
94ef460
5a2e9d4
77ed406
204d759
77ed406
 
5a2e9d4
204d759
 
94ef460
4f30320
b293ab8
77ed406
204d759
 
77ed406
204d759
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
5a2e9d4
204d759
 
 
 
b293ab8
204d759
 
102288d
b293ab8
 
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
# app.py - Robust Local Ollama Chatbot
# Features: Dynamic Connection, Model Selection, Robust History
import gradio as gr
import requests
import logging
import sys
import shutil
from datetime import datetime

# Logging
logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s | %(levelname)-8s | %(message)s",
    handlers=[logging.StreamHandler(sys.stdout)]
)
logger = logging.getLogger("CHATBOT")

OLLAMA_URL = "http://localhost:11434"

# Global state
current_model = ""
available_models = []

def get_ollama_models():
    """Fetch available models from Ollama."""
    try:
        r = requests.get(f"{OLLAMA_URL}/api/tags", timeout=2)
        if r.status_code == 200:
            return [m["name"] for m in r.json().get("models", [])]
    except:
        pass
    return []

def check_connection():
    """Check connection and update model list."""
    global current_model, available_models
    
    if shutil.which("ollama") is None:
        return "❌ **Ollama not installed**", gr.update(choices=[], value=None)
    
    models = get_ollama_models()
    available_models = models
    
    if not models:
        return "⚠️ Ollama running but **NO models found**", gr.update(choices=[], value=None)
    
    # Default logic: Prefer llama3.2, then deepseek, then first available
    if not current_model or current_model not in models:
        if "llama3.2:3b" in models:
            current_model = "llama3.2:3b"
        elif "deepseek-coder:6.7b-instruct-q6_K" in models:
            current_model = "deepseek-coder:6.7b-instruct-q6_K"
        else:
            current_model = models[0]
            
    return f"✅ Connected", gr.update(choices=models, value=current_model)

def chat_fn(message, history, model_name):
    """Chat function handling history and model selection."""
    if not model_name:
        return "Error: No model selected. Please check connection."
        
    # Build prompt
    prompt = "You are a helpful AI assistant.\n\n"
    for item in history:
        if isinstance(item, dict):
            role = item.get("role", "")
            content = item.get("content", "")
            if role == "user":
                prompt += f"User: {content}\n"
            elif role == "assistant":
                prompt += f"Assistant: {content}\n"
    prompt += f"User: {message}\nAssistant:"
    
    try:
        r = requests.post(
            f"{OLLAMA_URL}/api/generate",
            json={"model": model_name, "prompt": prompt, "stream": False},
            timeout=120
        )
        if r.status_code == 200:
            return r.json().get("response", "").strip()
        return f"Error: Ollama returned {r.status_code}"
    except Exception as e:
        return f"Error: {e}"

# UI
with gr.Blocks(title="Local Chatbot") as demo:
    gr.Markdown("# 🤖 Local AI Chatbot")
    
    with gr.Row():
        status_display = gr.Markdown("**Checking connection...**")
        refresh_btn = gr.Button("🔄 Check Connection", size="sm")
    
    with gr.Row():
        model_dropdown = gr.Dropdown(
            label="Select Model",
            choices=[],
            value=None,
            interactive=True
        )
    
    chat_interface = gr.ChatInterface(
        fn=chat_fn,
        additional_inputs=[model_dropdown],
    )
    
    with gr.Accordion("🛠️ Setup Instructions", open=False):
        gr.Markdown("""
        1. **Install Ollama:** [ollama.com](https://ollama.com)
        2. **Start Server:** `ollama serve`
        3. **Download Model:** `ollama pull llama3.2:3b`
        """)

    # Wiring
    def update_all():
        status, model_update = check_connection()
        return status, model_update

    refresh_btn.click(update_all, outputs=[status_display, model_dropdown])
    demo.load(update_all, outputs=[status_display, model_dropdown])

if __name__ == "__main__":
    demo.launch()