mirror of
https://gh.wpcy.net/https://github.com/discourse/discourse.git
synced 2026-06-19 10:44:14 +08:00
ab51aa7d81
Adds a `crypto` API to the JavaScript tool runner exposing HMAC
(SHA1/SHA256), hashing (MD5/SHA1/SHA256), RSA PKCS1v15 signing, base64
and base64url encoding, and cryptographically secure random bytes. All
functions are synchronous bridges to Ruby's OpenSSL with string or
Uint8Array inputs and a 10MB per-call limit, enabling webhook signature
verification, JWT signing for service accounts, and similar use cases.
Also splits the monolithic `tool_runner.rb` into per-feature modules
(`http`, `llm`, `index`, `upload`, `discourse`, `crypto`) with matching
spec files, and documents the discourse API's security model in the
tool preamble so authors understand the privilege boundaries when
non-admins trigger their tools.
New interface allows for stuff like this in a tool:
```
function signJWT(cred) {
const now = Math.floor(Date.now() / 1000);
const hdr = crypto.base64UrlEncode(JSON.stringify({ alg: "RS256", typ: "JWT" }));
const pay = crypto.base64UrlEncode(JSON.stringify({
iss: cred.client_email,
scope: "https://www.googleapis.com/auth/bigquery.readonly",
aud: "https://oauth2.googleapis.com/token",
iat: now,
exp: now + 3600,
}));
const msg = hdr + "." + pay;
return msg + "." + crypto.base64UrlEncode(crypto.signRsaSha256(cred.private_key, msg));
}
function getToken(cred) {
const jwt = signJWT(cred);
const r = http.post("https://oauth2.googleapis.com/token", {
headers: { "Content-Type": "application/x-www-form-urlencoded" },
body:
"grant_type=urn%3Aietf%3Aparams%3Aoauth%3Agrant-type%3Ajwt-bearer&assertion=" + jwt,
});
if (r.status !== 200) throw new Error("OAuth2 failed (" + r.status + "): " + r.body);
return JSON.parse(r.body).access_token;
}
function invoke(params) {
const cred = JSON.parse(secrets.get("json_cred"));
const token = getToken(cred);
const project = cred.project_id;
const url =
"https://bigquery.googleapis.com/bigquery/v2/projects/" +
encodeURIComponent(project) +
"/datasets";
const r = http.get(url, { headers: { Authorization: "Bearer " + token } });
if (r.status !== 200) throw new Error("BigQuery (" + r.status + "): " + r.body);
const body = JSON.parse(r.body);
const datasets = (body.datasets || []).map(function (ds) {
return ds.datasetReference.datasetId;
});
return { project: project, datasets: datasets };
}
function details() {
return "Lists all BigQuery datasets in the project";
}
```
Previously to achieve the same thing we would need to implement crypto
direct in JS
97 lines
2.7 KiB
Ruby
Vendored
97 lines
2.7 KiB
Ruby
Vendored
# frozen_string_literal: true
|
|
|
|
require "rails_helper"
|
|
|
|
RSpec.describe DiscourseAi::Agents::ToolRunner do
|
|
fab!(:llm_model) { Fabricate(:llm_model, name: "claude-2") }
|
|
let(:llm) { DiscourseAi::Completions::Llm.proxy(llm_model) }
|
|
fab!(:bot_user) { Discourse.system_user }
|
|
|
|
def create_tool(script:)
|
|
AiTool.create!(
|
|
name: "test #{SecureRandom.uuid}",
|
|
tool_name: "test_#{SecureRandom.uuid.underscore}",
|
|
description: "test",
|
|
parameters: [{ name: "query", type: "string", description: "perform a search" }],
|
|
script: script,
|
|
created_by_id: 1,
|
|
summary: "Test tool summary",
|
|
)
|
|
end
|
|
|
|
before { enable_current_plugin }
|
|
|
|
describe "LLM operations" do
|
|
it "has access to llm truncation tools" do
|
|
script = <<~JS
|
|
function invoke(params) {
|
|
return llm.truncate("Hello World", 1);
|
|
}
|
|
JS
|
|
|
|
tool = create_tool(script: script)
|
|
|
|
runner = tool.runner({}, llm: llm, bot_user: nil)
|
|
result = runner.invoke
|
|
|
|
expect(result).to eq("Hello")
|
|
end
|
|
|
|
it "is able to run llm completions" do
|
|
script = <<~JS
|
|
function invoke(params) {
|
|
return llm.generate("question two") + llm.generate(
|
|
{ messages: [
|
|
{ type: "system", content: "system message" },
|
|
{ type: "user", content: "user message" }
|
|
]}
|
|
);
|
|
}
|
|
JS
|
|
|
|
tool = create_tool(script: script)
|
|
|
|
result = nil
|
|
prompts = nil
|
|
responses = ["Hello ", "World"]
|
|
|
|
DiscourseAi::Completions::Llm.with_prepared_responses(responses) do |_, _, _prompts|
|
|
runner = tool.runner({}, llm: llm, bot_user: nil)
|
|
result = runner.invoke
|
|
prompts = _prompts
|
|
end
|
|
|
|
prompt =
|
|
DiscourseAi::Completions::Prompt.new(
|
|
"system message",
|
|
messages: [{ type: :user, content: "user message" }],
|
|
)
|
|
expect(result).to eq("Hello World")
|
|
expect(prompts[0]).to eq("question two")
|
|
expect(prompts[1]).to eq(prompt)
|
|
end
|
|
|
|
it "can generate JSON from LLM" do
|
|
tool_record =
|
|
AiTool.create!(
|
|
name: "test_tool",
|
|
tool_name: "test_tool",
|
|
description: "a test tool",
|
|
script: "function invoke() { return llm.generate('test', { json: true }); }",
|
|
summary: "test",
|
|
created_by_id: 1,
|
|
)
|
|
|
|
DiscourseAi::Completions::Llm.with_prepared_responses(
|
|
['{"key": "value"}'],
|
|
) do |_, _, _, prompt_options|
|
|
runner =
|
|
described_class.new(parameters: {}, llm: llm, bot_user: bot_user, tool: tool_record)
|
|
result = runner.invoke
|
|
expect(result).to eq({ "key" => "value" })
|
|
|
|
expect(prompt_options.last[:response_format]).to eq({ "type" => "json_object" })
|
|
end
|
|
end
|
|
end
|
|
end
|
