97 lines
3.8 KiB
Python
97 lines
3.8 KiB
Python
import requests
|
|
from config import Config
|
|
|
|
|
|
def list_ollama_models():
|
|
"""List available models on Ollama server"""
|
|
if not Config.OLLAMA_ENABLED:
|
|
return None, "Ollama is not enabled"
|
|
|
|
try:
|
|
url = f"{Config.OLLAMA_BASE_URL}/api/tags"
|
|
headers = {}
|
|
if Config.OLLAMA_API_KEY:
|
|
headers["Authorization"] = f"Bearer {Config.OLLAMA_API_KEY}"
|
|
|
|
response = requests.get(url, headers=headers, timeout=10)
|
|
response.raise_for_status()
|
|
|
|
result = response.json()
|
|
models = result.get('models', [])
|
|
model_names = [model.get('name', '') for model in models]
|
|
|
|
return model_names, None
|
|
except requests.exceptions.RequestException as e:
|
|
return None, f"Error listing models: {str(e)}"
|
|
except Exception as e:
|
|
return None, f"Unexpected error: {str(e)}"
|
|
|
|
|
|
def call_ollama(prompt, system_prompt=None):
|
|
"""Call Ollama API to generate text"""
|
|
if not Config.OLLAMA_ENABLED:
|
|
return None, "Ollama is not enabled"
|
|
|
|
try:
|
|
url = f"{Config.OLLAMA_BASE_URL}/api/generate"
|
|
payload = {
|
|
"model": Config.OLLAMA_MODEL,
|
|
"prompt": prompt,
|
|
"stream": False
|
|
}
|
|
|
|
if system_prompt:
|
|
payload["system"] = system_prompt
|
|
|
|
headers = {}
|
|
if Config.OLLAMA_API_KEY:
|
|
headers["Authorization"] = f"Bearer {Config.OLLAMA_API_KEY}"
|
|
|
|
print(f"Calling Ollama at {url} with model {Config.OLLAMA_MODEL}")
|
|
response = requests.post(url, json=payload, headers=headers, timeout=30)
|
|
response.raise_for_status()
|
|
|
|
result = response.json()
|
|
response_text = result.get('response', '').strip()
|
|
|
|
if not response_text:
|
|
return None, "Ollama returned empty response"
|
|
|
|
return response_text, None
|
|
except requests.exceptions.ConnectionError as e:
|
|
error_msg = f"Cannot connect to Ollama server at {Config.OLLAMA_BASE_URL}. Is Ollama running?"
|
|
print(f"Connection error: {error_msg}")
|
|
return None, error_msg
|
|
except requests.exceptions.Timeout:
|
|
error_msg = "Request to Ollama timed out after 30 seconds"
|
|
print(f"Timeout error: {error_msg}")
|
|
return None, error_msg
|
|
except requests.exceptions.HTTPError as e:
|
|
# Check if it's a model not found error
|
|
if e.response.status_code == 404:
|
|
try:
|
|
error_data = e.response.json()
|
|
if 'model' in error_data.get('error', '').lower() and 'not found' in error_data.get('error', '').lower():
|
|
# Try to get available models
|
|
available_models, _ = list_ollama_models()
|
|
if available_models:
|
|
error_msg = f"Model '{Config.OLLAMA_MODEL}' not found. Available models: {', '.join(available_models)}"
|
|
else:
|
|
error_msg = f"Model '{Config.OLLAMA_MODEL}' not found. Use 'ollama list' on the server to see available models."
|
|
else:
|
|
error_msg = f"HTTP error from Ollama: {e.response.status_code} - {e.response.text}"
|
|
except (ValueError, KeyError):
|
|
error_msg = f"HTTP error from Ollama: {e.response.status_code} - {e.response.text}"
|
|
else:
|
|
error_msg = f"HTTP error from Ollama: {e.response.status_code} - {e.response.text}"
|
|
print(f"HTTP error: {error_msg}")
|
|
return None, error_msg
|
|
except requests.exceptions.RequestException as e:
|
|
error_msg = f"Request error: {str(e)}"
|
|
print(f"Request error: {error_msg}")
|
|
return None, error_msg
|
|
except Exception as e:
|
|
error_msg = f"Unexpected error: {str(e)}"
|
|
print(f"Unexpected error: {error_msg}")
|
|
return None, error_msg
|