llmservice.registerAction

llmservice.registerAction

Registers custom server-side actions for extended functionality.

llmService.registerAction() is a server-side method provided by the Language Learning Model (LLM) Service. The function adds a custom action to the LLMService.

Syntax

llmService.registerAction(actionName, actionFunction);
llmService.registerAction(actionName, actionFunction);

Parameters

  • actionName - A string that defines the unique name for the action.
  • actionFunction - A function that describes the action. This function must be async and take an options object as a parameter. It should return a Promise that resolves to either a ReadableStream or an object.

Requirements

You can use the usellm "https://usellm.org/api/llm" service URL for testing. But if you're working on your personal project you'll need to use your own Replicate API. Go to the Replicate API page, create your own API token and paste the API Key in .env.local file (check .env.example). You'll also have to pass the API key saved in .env.local file while defining CreateLLMService.

Example

Below is an example of how to register a custom action. In this case, we are registering a replicateText action. The action makes a POST requrest to Replicate.com with a text and gets a Response with the output "Hello" + text (live demo):

import { createLLMService } from "usellm";
import { makeErrorResponse } from "@/usellm/shared/utils";
import { Ratelimit } from "@upstash/ratelimit";
import { Redis } from "@upstash/redis";
 
const llmService = createLLMService({
  fetcher: fetch,
  actions: ["generateHighResImage"],
  isAllowed: async () => {
    // check if rate limiting has been set up using Upstash Redis REST API
    if (
      process.env.UPSTASH_REDIS_REST_URL &&
      process.env.UPSTASH_REDIS_REST_TOKEN
    ) {
      const { success } = await ratelimit.limit("api");
      return success;
    }
    return true;
  },
});
 
// Register new action
llmService.registerAction("generateHighResImage", async (options: any) => {
  const { prompt } = options;
  if (!prompt) {
    throw makeErrorResponse("'prompt' is required", 400);
  }
 
  const REPLICATE_API_URL = "https://api.replicate.com/v1/predictions";
  const STABLE_DIFFUSION_MODEL_ID =
    "db21e45d3f7023abc2a46ee38a23973f6dce16bb082a930b0c49861f96d1e5bf";
 
  const response1 = await fetch(REPLICATE_API_URL, {
    method: "POST",
    headers: {
      "Content-Type": "application/json",
      Authorization: `Token ${process.env.REPLICATE_API_TOKEN}`,
    },
    body: JSON.stringify({
      version: STABLE_DIFFUSION_MODEL_ID,
      input: {
        prompt: prompt,
      },
    }),
  });
 
  if (!response1.ok) {
    throw new Error(await response1.text());
  }
  const { id: prediction_id } = await response1.json();
 
  // Wait for 30 seconds to run the model
  const sleep = async (milliseconds: number) => {
    await new Promise((resolve) => {
      return setTimeout(resolve, milliseconds);
    });
  };
  await sleep(30000);
 
  // Get the model response from Replicate
  const link = REPLICATE_API_URL + "/" + prediction_id;
 
  const response2 = await fetch(link, {
    method: "GET",
    headers: {
      Authorization: `Token ${process.env.REPLICATE_API_TOKEN}`,
    },
  });
 
  const data = await response2.json();
  return { data };
});
import { createLLMService } from "usellm";
import { makeErrorResponse } from "@/usellm/shared/utils";
import { Ratelimit } from "@upstash/ratelimit";
import { Redis } from "@upstash/redis";
 
const llmService = createLLMService({
  fetcher: fetch,
  actions: ["generateHighResImage"],
  isAllowed: async () => {
    // check if rate limiting has been set up using Upstash Redis REST API
    if (
      process.env.UPSTASH_REDIS_REST_URL &&
      process.env.UPSTASH_REDIS_REST_TOKEN
    ) {
      const { success } = await ratelimit.limit("api");
      return success;
    }
    return true;
  },
});
 
// Register new action
llmService.registerAction("generateHighResImage", async (options: any) => {
  const { prompt } = options;
  if (!prompt) {
    throw makeErrorResponse("'prompt' is required", 400);
  }
 
  const REPLICATE_API_URL = "https://api.replicate.com/v1/predictions";
  const STABLE_DIFFUSION_MODEL_ID =
    "db21e45d3f7023abc2a46ee38a23973f6dce16bb082a930b0c49861f96d1e5bf";
 
  const response1 = await fetch(REPLICATE_API_URL, {
    method: "POST",
    headers: {
      "Content-Type": "application/json",
      Authorization: `Token ${process.env.REPLICATE_API_TOKEN}`,
    },
    body: JSON.stringify({
      version: STABLE_DIFFUSION_MODEL_ID,
      input: {
        prompt: prompt,
      },
    }),
  });
 
  if (!response1.ok) {
    throw new Error(await response1.text());
  }
  const { id: prediction_id } = await response1.json();
 
  // Wait for 30 seconds to run the model
  const sleep = async (milliseconds: number) => {
    await new Promise((resolve) => {
      return setTimeout(resolve, milliseconds);
    });
  };
  await sleep(30000);
 
  // Get the model response from Replicate
  const link = REPLICATE_API_URL + "/" + prediction_id;
 
  const response2 = await fetch(link, {
    method: "GET",
    headers: {
      Authorization: `Token ${process.env.REPLICATE_API_TOKEN}`,
    },
  });
 
  const data = await response2.json();
  return { data };
});

In this example, the generateHighResImage action takes a prompt input, makes a POST request to a Replicate model, waits for the model to generate the high-resolution image, then retrieves and returns the image. If the prompt input is not provided, an error is thrown.

To know how to call the registered action check the llm.callAction page.

By using llmService.registerAction(), you can seamlessly integrate custom actions into your LLMService, tailoring its functionality to your specific needs.