| |
| """ |
| Model Manager for Police Bot AI Runtime |
| Helps manage and switch between different Ollama models |
| """ |
|
|
| import os |
| import sys |
| import requests |
| import json |
| from config import config |
|
|
| def print_banner(): |
| """Print banner""" |
| print("=" * 60) |
| print("🤖 Police Bot AI - Model Manager") |
| print("=" * 60) |
|
|
| def check_ollama_status(): |
| """Check if Ollama is running and get available models""" |
| try: |
| response = requests.get("http://localhost:11434/api/tags", timeout=5) |
| if response.status_code == 200: |
| data = response.json() |
| models = data.get("models", []) |
| return True, models |
| else: |
| return False, [] |
| except: |
| return False, [] |
|
|
| def list_available_models(): |
| """List all available models in Ollama""" |
| print("\n📋 Available Models in Ollama:") |
| print("-" * 40) |
| |
| ollama_running, models = check_ollama_status() |
| |
| if not ollama_running: |
| print("❌ Ollama is not running") |
| print("Please start Ollama first: ollama serve") |
| return [] |
| |
| if not models: |
| print("⚠️ No models found in Ollama") |
| print("You can pull models using:") |
| print(" ollama pull mistral") |
| print(" ollama pull llama3") |
| print(" ollama pull police-bot") |
| return [] |
| |
| available_models = [] |
| for i, model in enumerate(models, 1): |
| model_name = model.get("name", "Unknown") |
| model_size = model.get("size", 0) |
| size_mb = model_size / (1024 * 1024) if model_size > 0 else 0 |
| |
| print(f"{i:2d}. {model_name:<20} ({size_mb:.0f} MB)") |
| available_models.append(model_name) |
| |
| return available_models |
|
|
| def list_configured_models(): |
| """List models configured in the system""" |
| print("\n⚙️ Configured Models in System:") |
| print("-" * 40) |
| |
| configured_models = config.list_available_models() |
| current_model = config.OLLAMA_MODEL |
| |
| for i, model in enumerate(configured_models, 1): |
| status = "✓" if model == current_model else " " |
| print(f"{i:2d}. [{status}] {model}") |
| |
| return configured_models |
|
|
| def switch_model(): |
| """Switch to a different model""" |
| print("\n🔄 Switch Model") |
| print("-" * 40) |
| |
| configured_models = config.list_available_models() |
| current_model = config.OLLAMA_MODEL |
| |
| print(f"Current model: {current_model}") |
| print("\nAvailable models:") |
| |
| for i, model in enumerate(configured_models, 1): |
| print(f"{i}. {model}") |
| |
| try: |
| choice = input(f"\nEnter choice (1-{len(configured_models)}): ").strip() |
| if choice: |
| model_index = int(choice) - 1 |
| if 0 <= model_index < len(configured_models): |
| selected_model = configured_models[model_index] |
| |
| |
| ollama_running, models = check_ollama_status() |
| if ollama_running: |
| model_names = [m.get("name") for m in models] |
| if selected_model not in model_names: |
| print(f"⚠️ Warning: {selected_model} is not available in Ollama") |
| pull = input("Would you like to pull it? (y/n): ").lower().strip() |
| if pull == 'y': |
| print(f"Pulling {selected_model}...") |
| os.system(f"ollama pull {selected_model}") |
| |
| |
| if config.switch_model(selected_model): |
| print(f"✅ Switched to {selected_model}") |
| |
| |
| save_model_to_env(selected_model) |
| else: |
| print(f"❌ Failed to switch to {selected_model}") |
| else: |
| print("❌ Invalid choice") |
| else: |
| print("❌ No choice made") |
| except ValueError: |
| print("❌ Invalid input") |
|
|
| def save_model_to_env(model_name): |
| """Save model selection to .env file""" |
| env_file = ".env" |
| env_content = f"OLLAMA_MODEL={model_name}\n" |
| |
| try: |
| with open(env_file, 'w') as f: |
| f.write(env_content) |
| print(f"✅ Model selection saved to {env_file}") |
| except Exception as e: |
| print(f"⚠️ Could not save to {env_file}: {e}") |
|
|
| def test_model(): |
| """Test the current model""" |
| print("\n🧪 Test Current Model") |
| print("-" * 40) |
| |
| current_model = config.OLLAMA_MODEL |
| print(f"Testing model: {current_model}") |
| |
| |
| ollama_running, models = check_ollama_status() |
| if not ollama_running: |
| print("❌ Ollama is not running") |
| return |
| |
| |
| model_names = [m.get("name") for m in models] |
| if current_model not in model_names: |
| print(f"❌ {current_model} is not available in Ollama") |
| print(f"Available models: {', '.join(model_names)}") |
| return |
| |
| |
| try: |
| from police_runtime import PoliceBotRuntime |
| |
| runtime = PoliceBotRuntime() |
| test_prompt = "Hello, how are you today?" |
| |
| print(f"Sending test prompt: '{test_prompt}'") |
| response = runtime.get_llama3_response(test_prompt) |
| |
| if response: |
| print(f"✅ Model responded: {response[:100]}...") |
| else: |
| print("❌ No response from model") |
| |
| except Exception as e: |
| print(f"❌ Error testing model: {e}") |
|
|
| def show_model_info(): |
| """Show information about the current model""" |
| print("\nℹ️ Current Model Information") |
| print("-" * 40) |
| |
| current_model = config.OLLAMA_MODEL |
| model_config = config.get_current_model_config() |
| |
| print(f"Model: {current_model}") |
| print(f"Temperature: {model_config['temperature']}") |
| print(f"Top P: {model_config['top_p']}") |
| print(f"Max Tokens: {model_config['max_tokens']}") |
| |
| print(f"\nSystem Prompt Preview:") |
| print("-" * 20) |
| prompt_preview = model_config['system_prompt'][:200] + "..." if len(model_config['system_prompt']) > 200 else model_config['system_prompt'] |
| print(prompt_preview) |
|
|
| def pull_model(): |
| """Pull a model from Ollama""" |
| print("\n📥 Pull Model from Ollama") |
| print("-" * 40) |
| |
| model_name = input("Enter model name to pull (e.g., mistral, llama3): ").strip() |
| if model_name: |
| print(f"Pulling {model_name}...") |
| result = os.system(f"ollama pull {model_name}") |
| |
| if result == 0: |
| print(f"✅ Successfully pulled {model_name}") |
| else: |
| print(f"❌ Failed to pull {model_name}") |
| else: |
| print("❌ No model name provided") |
|
|
| def main(): |
| """Main function""" |
| print_banner() |
| |
| while True: |
| print("\nOptions:") |
| print("1. List available models in Ollama") |
| print("2. List configured models") |
| print("3. Switch model") |
| print("4. Test current model") |
| print("5. Show model information") |
| print("6. Pull model from Ollama") |
| print("7. Exit") |
| |
| try: |
| choice = input("\nEnter choice (1-7): ").strip() |
| |
| if choice == '1': |
| list_available_models() |
| elif choice == '2': |
| list_configured_models() |
| elif choice == '3': |
| switch_model() |
| elif choice == '4': |
| test_model() |
| elif choice == '5': |
| show_model_info() |
| elif choice == '6': |
| pull_model() |
| elif choice == '7': |
| print("Goodbye! 👋") |
| break |
| else: |
| print("❌ Invalid choice") |
| |
| except KeyboardInterrupt: |
| print("\nGoodbye! 👋") |
| break |
| except Exception as e: |
| print(f"❌ Error: {e}") |
|
|
| if __name__ == "__main__": |
| main() |
