discourse-ai/spec/shared/inference/openai_embeddings_spec.rb

66 lines
1.9 KiB
Ruby

# frozen_string_literal: true
describe DiscourseAi::Inference::OpenAiEmbeddings do
it "supports azure embeddings" do
SiteSetting.ai_openai_embeddings_url =
"https://my-company.openai.azure.com/openai/deployments/embeddings-deployment/embeddings?api-version=2023-05-15"
SiteSetting.ai_openai_api_key = "123456"
body_json = {
usage: {
prompt_tokens: 1,
total_tokens: 1,
},
data: [{ object: "embedding", embedding: [0.0, 0.1] }],
}.to_json
stub_request(
:post,
"https://my-company.openai.azure.com/openai/deployments/embeddings-deployment/embeddings?api-version=2023-05-15",
).with(
body: "{\"model\":\"text-embedding-ada-002\",\"input\":\"hello\"}",
headers: {
"Api-Key" => "123456",
"Content-Type" => "application/json",
},
).to_return(status: 200, body: body_json, headers: {})
result =
DiscourseAi::Inference::OpenAiEmbeddings.instance(model: "text-embedding-ada-002").perform!(
"hello",
)
expect(result).to eq([0.0, 0.1])
end
it "supports openai embeddings" do
SiteSetting.ai_openai_api_key = "123456"
body_json = {
usage: {
prompt_tokens: 1,
total_tokens: 1,
},
data: [{ object: "embedding", embedding: [0.0, 0.1] }],
}.to_json
body = { model: "text-embedding-ada-002", input: "hello", dimensions: 1000 }.to_json
stub_request(:post, "https://api.openai.com/v1/embeddings").with(
body: body,
headers: {
"Authorization" => "Bearer 123456",
"Content-Type" => "application/json",
},
).to_return(status: 200, body: body_json, headers: {})
result =
DiscourseAi::Inference::OpenAiEmbeddings.instance(
model: "text-embedding-ada-002",
dimensions: 1000,
).perform!("hello")
expect(result).to eq([0.0, 0.1])
end
end