import os

import pytest
import httpx
from httpx import HTTPError, ConnectError

from chromadb.utils.embedding_functions import OllamaEmbeddingFunction


def test_ollama() -> None:
    """
    To set up the Ollama server, follow instructions at: https://github.com/ollama/ollama?tab=readme-ov-file
    Export the OLLAMA_SERVER_URL and OLLAMA_MODEL environment variables.
    """
    if (
        os.environ.get("OLLAMA_SERVER_URL") is None
        or os.environ.get("OLLAMA_MODEL") is None
    ):
        pytest.skip(
            "OLLAMA_SERVER_URL or OLLAMA_MODEL environment variable not set. Skipping test."
        )
    try:
        response = httpx.get(os.environ.get("OLLAMA_SERVER_URL", ""))
        # If the response was successful, no Exception will be raised
        response.raise_for_status()
    except (HTTPError, ConnectError):
        pytest.skip("Ollama server not running. Skipping test.")
    ef = OllamaEmbeddingFunction(
        model_name=os.environ.get("OLLAMA_MODEL") or "nomic-embed-text",
        url=f"{os.environ.get('OLLAMA_SERVER_URL')}/embeddings",
    )
    embeddings = ef(["Here is an article about llamas...", "this is another article"])
    assert len(embeddings) == 2