summaryrefslogtreecommitdiffstats
path: root/krebs/5pkgs/simple/vicuna-chat/default.nix
blob: db15899d6778d8e9c0af96572713417e77c57d97 (plain)
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
{ pkgs, ... }:
pkgs.writers.writeDashBin "vicuna-chat" ''
  set -efu

  export PATH=${with pkgs; lib.makeBinPath [
    coreutils
    curl
    jq
  ]}

  CONTEXT=''${CONTEXT:-$(date -Id)}
  PROMPT=$*

  if ! test -e "$CONTEXT"; then
    echo -n 'null' > "$CONTEXT"
  fi

  add_to_context() {
    jq -rc --argjson message "$1" '. + [$message]' "$CONTEXT" > "$CONTEXT.tmp"
    mv "$CONTEXT.tmp" "$CONTEXT"
  }

  add_to_context "{\"role\": \"user\", \"content\": \"$PROMPT\"}"
  response=$(
    jq -nc --slurpfile context "$CONTEXT" '{
      model: "vicuna-13b-v1.5-16k",
      messages: $context[0],
    }' |
      curl -Ss http://vicuna.r/v1/chat/completions -H 'Content-Type: application/json' -d @-
  )
  add_to_context "$(jq -rcn --argjson response "$response" '$response.choices[0].message')"
  jq -rcn --argjson response "$response" '$response.choices[0].message.content'
''