ollama

1
import json
2
import requests
3

4
# NOTE: ollama must be running for this to work, start the ollama app or run `ollama serve`
5
model = "llama2"  # TODO: update this for whatever model you wish to use
6

7

8
def chat(messages):
9
    r = requests.post(
10
        "http://0.0.0.0:11434/api/chat",
11
        json={"model": model, "messages": messages, "stream": True},
12
    )
13
    r.raise_for_status()
14
    output = ""
15

16
    for line in r.iter_lines():
17
        body = json.loads(line)
18
        if "error" in body:
19
            raise Exception(body["error"])
20
        if body.get("done") is False:
21
            message = body.get("message", "")
22
            content = message.get("content", "")
23
            output += content
24
            # the response streams one token at a time, print that as we receive it
25
            print(content, end="", flush=True)
26

27
        if body.get("done", False):
28
            message["content"] = output
29
            return message
30

31

32
def main():
33
    messages = []
34

35
    while True:
36
        user_input = input("Enter a prompt: ")
37
        if not user_input:
38
            exit()
39
        print()
40
        messages.append({"role": "user", "content": user_input})
41
        message = chat(messages)
42
        messages.append(message)
43
        print("\n\n")
44

45

46
if __name__ == "__main__":
47
    main()
48