Query 🤗TGI /info endpoint to get model name and context window (#755)

* query TGI /info endpoint for model name and context length * remove commented-out line * add parseInt --------- Co-authored-by: Simon Müller <[email protected]>
continuedev · Jan 16, 2024 · 0c77746 · 0c77746
1 parent dbbccef
commit 0c77746
Showing 1 changed file with 19 additions and 0 deletions.
diff --git a/core/llm/llms/HuggingFaceTGI.ts b/core/llm/llms/HuggingFaceTGI.ts
@@ -8,6 +8,25 @@ class HuggingFaceTGI extends BaseLLM {
     apiBase: "http://localhost:8080",
   };
 
+  constructor(options: LLMOptions) {
+    super(options);
+
+    this.fetch(`${this.apiBase}/info`, {
+      method: "GET",
+    }).then(async (response) => {
+      if (response.status !== 200) {
+        console.warn(
+          "Error calling Hugging Face TGI /info endpoint: ",
+          await response.text()
+        );
+        return;
+      }
+      const json = await response.json();
+      this.model = json.model_id;
+      this.contextLength = parseInt(json.max_input_length);
+    });
+  }
+
   private _convertArgs(options: CompletionOptions, prompt: string) {
     const finalOptions = {
       max_new_tokens: options.maxTokens,