OnionSpy/model.py at main · CoderRony955/OnionSpy · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
from rich.console import Console
from ollama import chat, list
import subprocess
from config import (
    ollama_cloud_model,
    ollama_api_key,
    ollama_local_model
)
import random
import sys


# logger
from logger import logger

console = Console()


def check_local_models():
    try:
        models = list()
        if not models["models"]:
            return False
        return True
    except Exception as e:
        logger.error(e)
        return


def ask(prompt: str):
    try:
        # check if ollama installed in machine or not
        command = subprocess.run(
            ["ollama", "--version"], shell=True, capture_output=True)

        # if ollama is not install in machine
        if command.stderr:
            logger.error("don't have installed Ollama")
            return False

        # if doesn't have installed local models but have cloud api key and cloud model access
        if not check_local_models() and ollama_api_key and ollama_cloud_model:
            response = chat(
                model=ollama_cloud_model,
                messages=[
                    {'role': 'user', 'content': prompt},
                ],
            )
            return response['message']['content']

        # if doesn't have cloud api and cloud api access but have installed local models
        elif not ollama_api_key and not ollama_cloud_model and check_local_models():
            response = chat(
                model=ollama_local_model,
                messages=[
                    {'role': 'user', 'content': prompt},
                ],
            )
            return response['message']['content']

        # elif have both then choose cloud model for better and fast response
        elif check_local_models() and ollama_api_key and ollama_cloud_model:
            response = chat(
                model=ollama_cloud_model,
                messages=[
                    {'role': 'user', 'content': prompt},
                ],
            )
            return response['message']['content']

        # if doesn't have both
        else:
            return False
    except KeyboardInterrupt:
        from operations.keywords import Keywords
        console.print(
            f"\n[bold green]{random.choice(Keywords.goodbye_keywords)}![/bold green]")
        sys.exit(1)

    except Exception as e:
        logger.error(e)
        console.print(
            f"[bold red]Something went wrong! Please check the logs for more details.[/bold red]")
        return False