discourse-ai/spec/support/embeddings_generation_stubs.rb

24 lines
815 B
Ruby

# frozen_string_literal: true
class EmbeddingsGenerationStubs
class << self
def discourse_service(model, string, embedding)
model = "bge-large-en-v1.5" if model == "bge-large-en"
WebMock
.stub_request(
:post,
"#{SiteSetting.ai_embeddings_discourse_service_api_endpoint}/api/v1/classify",
)
.with(body: JSON.dump({ model: model, content: string }))
.to_return(status: 200, body: JSON.dump(embedding))
end
def openai_service(model, string, embedding, extra_args: {})
WebMock
.stub_request(:post, "https://api.openai.com/v1/embeddings")
.with(body: JSON.dump({ model: model, input: string }.merge(extra_args)))
.to_return(status: 200, body: JSON.dump({ data: [{ embedding: embedding }] }))
end
end
end