completion.js 5.8 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204
  1. const paramDefaults = {
  2. stream: true,
  3. n_predict: 500,
  4. temperature: 0.2,
  5. stop: ["</s>"]
  6. };
  7. let generation_settings = null;
  8. // Completes the prompt as a generator. Recommended for most use cases.
  9. //
  10. // Example:
  11. //
  12. // import { llama } from '/completion.js'
  13. //
  14. // const request = llama("Tell me a joke", {n_predict: 800})
  15. // for await (const chunk of request) {
  16. // document.write(chunk.data.content)
  17. // }
  18. //
  19. export async function* llama(prompt, params = {}, config = {}) {
  20. let controller = config.controller;
  21. const api_url = config.api_url || "";
  22. if (!controller) {
  23. controller = new AbortController();
  24. }
  25. const completionParams = { ...paramDefaults, ...params, prompt };
  26. const response = await fetch(`${api_url}/completion`, {
  27. method: 'POST',
  28. body: JSON.stringify(completionParams),
  29. headers: {
  30. 'Connection': 'keep-alive',
  31. 'Content-Type': 'application/json',
  32. 'Accept': 'text/event-stream',
  33. ...(params.api_key ? {'Authorization': `Bearer ${params.api_key}`} : {})
  34. },
  35. signal: controller.signal,
  36. });
  37. const reader = response.body.getReader();
  38. const decoder = new TextDecoder();
  39. let content = "";
  40. let leftover = ""; // Buffer for partially read lines
  41. try {
  42. let cont = true;
  43. while (cont) {
  44. const result = await reader.read();
  45. if (result.done) {
  46. break;
  47. }
  48. // Add any leftover data to the current chunk of data
  49. const text = leftover + decoder.decode(result.value);
  50. // Check if the last character is a line break
  51. const endsWithLineBreak = text.endsWith('\n');
  52. // Split the text into lines
  53. let lines = text.split('\n');
  54. // If the text doesn't end with a line break, then the last line is incomplete
  55. // Store it in leftover to be added to the next chunk of data
  56. if (!endsWithLineBreak) {
  57. leftover = lines.pop();
  58. } else {
  59. leftover = ""; // Reset leftover if we have a line break at the end
  60. }
  61. // Parse all sse events and add them to result
  62. const regex = /^(\S+):\s(.*)$/gm;
  63. for (const line of lines) {
  64. const match = regex.exec(line);
  65. if (match) {
  66. result[match[1]] = match[2]
  67. // since we know this is llama.cpp, let's just decode the json in data
  68. if (result.data) {
  69. result.data = JSON.parse(result.data);
  70. content += result.data.content;
  71. // yield
  72. yield result;
  73. // if we got a stop token from server, we will break here
  74. if (result.data.stop) {
  75. if (result.data.generation_settings) {
  76. generation_settings = result.data.generation_settings;
  77. }
  78. cont = false;
  79. break;
  80. }
  81. }
  82. if (result.error) {
  83. try {
  84. result.error = JSON.parse(result.error);
  85. if (result.error.message.includes('slot unavailable')) {
  86. // Throw an error to be caught by upstream callers
  87. throw new Error('slot unavailable');
  88. } else {
  89. console.error(`llama.cpp error [${result.error.code} - ${result.error.type}]: ${result.error.message}`);
  90. }
  91. } catch(e) {
  92. console.error(`llama.cpp error ${result.error}`)
  93. }
  94. }
  95. }
  96. }
  97. }
  98. } catch (e) {
  99. if (e.name !== 'AbortError') {
  100. console.error("llama error: ", e);
  101. }
  102. throw e;
  103. }
  104. finally {
  105. controller.abort();
  106. }
  107. return content;
  108. }
  109. // Call llama, return an event target that you can subscribe to
  110. //
  111. // Example:
  112. //
  113. // import { llamaEventTarget } from '/completion.js'
  114. //
  115. // const conn = llamaEventTarget(prompt)
  116. // conn.addEventListener("message", (chunk) => {
  117. // document.write(chunk.detail.content)
  118. // })
  119. //
  120. export const llamaEventTarget = (prompt, params = {}, config = {}) => {
  121. const eventTarget = new EventTarget();
  122. (async () => {
  123. let content = "";
  124. for await (const chunk of llama(prompt, params, config)) {
  125. if (chunk.data) {
  126. content += chunk.data.content;
  127. eventTarget.dispatchEvent(new CustomEvent("message", { detail: chunk.data }));
  128. }
  129. if (chunk.data.generation_settings) {
  130. eventTarget.dispatchEvent(new CustomEvent("generation_settings", { detail: chunk.data.generation_settings }));
  131. }
  132. if (chunk.data.timings) {
  133. eventTarget.dispatchEvent(new CustomEvent("timings", { detail: chunk.data.timings }));
  134. }
  135. }
  136. eventTarget.dispatchEvent(new CustomEvent("done", { detail: { content } }));
  137. })();
  138. return eventTarget;
  139. }
  140. // Call llama, return a promise that resolves to the completed text. This does not support streaming
  141. //
  142. // Example:
  143. //
  144. // llamaPromise(prompt).then((content) => {
  145. // document.write(content)
  146. // })
  147. //
  148. // or
  149. //
  150. // const content = await llamaPromise(prompt)
  151. // document.write(content)
  152. //
  153. export const llamaPromise = (prompt, params = {}, config = {}) => {
  154. return new Promise(async (resolve, reject) => {
  155. let content = "";
  156. try {
  157. for await (const chunk of llama(prompt, params, config)) {
  158. content += chunk.data.content;
  159. }
  160. resolve(content);
  161. } catch (error) {
  162. reject(error);
  163. }
  164. });
  165. };
  166. /**
  167. * (deprecated)
  168. */
  169. export const llamaComplete = async (params, controller, callback) => {
  170. for await (const chunk of llama(params.prompt, params, { controller })) {
  171. callback(chunk);
  172. }
  173. }
  174. // Get the model info from the server. This is useful for getting the context window and so on.
  175. export const llamaModelInfo = async (config = {}) => {
  176. if (!generation_settings) {
  177. const api_url = config.api_url || "";
  178. const props = await fetch(`${api_url}/props`).then(r => r.json());
  179. generation_settings = props.default_generation_settings;
  180. }
  181. return generation_settings;
  182. }