discourse-ai/spec/shared/inference/openai_embeddings_spec.rb
Sam d1ab79e82f
FEATURE: Add Azure cognitive service support (#93)
The new site settings:

ai_openai_gpt35_url : distribution for GPT 16k
ai_openai_gpt4_url: distribution for GPT 4
ai_openai_embeddings_url: distribution for ada2

If untouched we will simply use OpenAI endpoints.

Azure requires 1 URL per model, OpenAI allows a single URL to serve multiple models. Hence the new settings.
2023-06-21 10:39:51 +10:00

60 lines
1.9 KiB
Ruby

# frozen_string_literal: true
require "rails_helper"
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.perform!("hello")
expect(result[:usage]).to eq({ prompt_tokens: 1, total_tokens: 1 })
expect(result[:data].first).to eq({ object: "embedding", embedding: [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
stub_request(:post, "https://api.openai.com/v1/embeddings").with(
body: "{\"model\":\"text-embedding-ada-002\",\"input\":\"hello\"}",
headers: {
"Authorization" => "Bearer 123456",
"Content-Type" => "application/json",
},
).to_return(status: 200, body: body_json, headers: {})
result = DiscourseAi::Inference::OpenAiEmbeddings.perform!("hello")
expect(result[:usage]).to eq({ prompt_tokens: 1, total_tokens: 1 })
expect(result[:data].first).to eq({ object: "embedding", embedding: [0.0, 0.1] })
end
end