import requests


class Ollama:
    def __init__(self, model="qwen2.5:14b", server_url="http://192.168.1.200:11434"):
        self.model = model
        self.server_url = server_url

    def is_service_running(self):
        """检查 Ollama 服务是否正在运行"""
        try:
            url = f"{self.server_url}/api/generate"
            headers = {"Content-Type": "application/json"}
            payload = {
                "model": self.model,
                "prompt": "ping",  # 使用轻量级的 prompt 测试
                "stream": False
            }
            output = requests.post(url, headers=headers, json=payload, timeout=60)
            if output.status_code == 200:
                return True
            else:
                return False
        except requests.exceptions.RequestException:
            return False

    def generate_text(self, input_message: str):
        url = f"{self.server_url}/api/generate"
        headers = {"Content-Type": "application/json"}
        payload = {
            "model": self.model,
            "prompt": input_message,
            "stream": False,  # 关闭流式传输
            "options": {
                "temperature": 0.7,
                "top_p": 0.9,
                "num_ctx": 20480
            }
        }
        try:
            output = requests.post(url, headers=headers, json=payload)
            output.raise_for_status()  # 检查HTTP请求是否成功
            result = output.json()
            return result.get("response", "No response from model.")
        except requests.exceptions.RequestException as e:
            return f"Error calling Ollama API: {e}"


if __name__ == "__main__":
    # 示例提示文本
    ollama = Ollama()
    print("Ollama 服务是否正在运行：", ollama.is_service_running())
    prompt = "请用中文描述机器学习的基本概念。"
    response = ollama.generate_text(prompt)
    print("模型生成的文本：")
    print(response)