1
/*
2
Multiresponse request/response conat server that makes
3
llm's available to users of CoCalc.
4

5
Query with a multiResponse request at this subject
6

7
    llm.account-{account_id}
8

9
We will also support llm.project-{project_id} at some point to make
10
it so projects can directly use llm's... but first we need to figure out
11
how paying for that would work.
12
*/
13

14
import { conat } from "@cocalc/conat/client";
15
import { isValidUUID } from "@cocalc/util/misc";
16
import type { Subscription } from "@cocalc/conat/core/client";
17

18
export const SUBJECT = process.env.COCALC_TEST_MODE ? "llm-test" : "llm";
19

20
export function llmSubject({
21
  account_id,
22
  project_id,
23
}: {
24
  account_id?: string;
25
  project_id?: string;
26
}) {
27
  if (account_id) {
28
    return `${SUBJECT}.account-${account_id}.api`;
29
  } else if (project_id) {
30
    return `${SUBJECT}.project-${project_id}.api`;
31
  } else {
32
    return `${SUBJECT}.hub.api`;
33
  }
34
}
35

36
function getUserId(subject: string): string {
37
  if (subject.startsWith(`${SUBJECT}.account-`)) {
38
    return subject.slice(
39
      `${SUBJECT}.account-`.length,
40
      `${SUBJECT}.account-`.length + 36,
41
    );
42
  }
43
  if (subject.startsWith(`${SUBJECT}.project-`)) {
44
    return subject.slice(
45
      `${SUBJECT}.project-`.length,
46
      `${SUBJECT}.project-`.length + 36,
47
    );
48
  }
49

50
  return "hub";
51
}
52

53
let sub: Subscription | null = null;
54
export async function init(evaluate) {
55
  const cn = await conat();
56
  sub = await cn.subscribe(`${SUBJECT}.*.api`, { queue: "q" });
57
  listen(evaluate);
58
}
59

60
export async function close() {
61
  if (sub == null) {
62
    return;
63
  }
64
  sub.drain();
65
  sub = null;
66
}
67

68
async function listen(evaluate) {
69
  if (sub == null) {
70
    throw Error("must init first");
71
  }
72
  for await (const mesg of sub) {
73
    handleMessage(mesg, evaluate);
74
  }
75
}
76

77
async function handleMessage(mesg, evaluate) {
78
  const options = mesg.data;
79

80
  let seq = 0;
81
  const respond = ({ text, error }: { text?: string; error?: string }) => {
82
    mesg.respondSync({ text, error, seq });
83
    seq += 1;
84
  };
85

86
  let done = false;
87
  const end = () => {
88
    if (done) return;
89
    done = true;
90
    // end response stream with null payload.
91
    mesg.respondSync(null);
92
  };
93

94
  const stream = (text?) => {
95
    if (done) return;
96
    if (text != null) {
97
      respond({ text });
98
    } else {
99
      end();
100
    }
101
  };
102

103
  try {
104
    // SECURITY: verify that the account_id claimed in options matches
105
    // with the subject the user published on (which proves they are who they say):
106
    // TODO: support project_id equally here.
107
    if (!isValidUUID(options.account_id)) {
108
      throw Error("account_id must be a valid uuid");
109
    }
110
    if (options.account_id != getUserId(mesg.subject)) {
111
      throw Error("account_id is invalid");
112
    }
113
    // NOTE: this evaluate does NOT await until all evaluation is done.
114
    // Instead evaluation is terminated by stream(undefined)!!
115
    await evaluate({ ...options, stream });
116
  } catch (err) {
117
    if (!done) {
118
      respond({ error: `${err}` });
119
      end();
120
    }
121
  }
122
}
123

124