# Ollama integration for NovaChat
from ollama import Client
import logging

def list_models() -> list:
    try:
        client = Client()
        return [model['name'] for model in client.list()['models']]
    except Exception as e:
        logging.error(f"Model listing failed: {str(e)}")
        return []


def pull_model(model_name: str) -> None:
    try:
        client = Client()
        client.pull(model_name, stream=False)
    except Exception as e:
        logging.error(f"Model pull failed: {str(e)}")
        raise
    response = requests.post(
        'http://localhost:11434/api/pull',
        json={'name': model_name},
        stream=True
    )
    
    for line in response.iter_lines():
        status = json.loads(line)
        if status.get('error'):
            raise Exception(status['error'])


def query_ollama(prompt, model="cogito", temperature=0.7, max_tokens=150):
    """Query local Ollama server using its REST API"""
    try:
        response = requests.post(
            'http://localhost:11434/api/generate',
            json={
                'model': model,
                'prompt': prompt,
                'temperature': temperature,
                'max_tokens': max_tokens,
                'stream': False
            },
            timeout=30
        )
        response.raise_for_status()
        return response.json().get('response', '')
    except Exception as e:
        return f"Ollama Error: {str(e)}"