Skip to content

Commit

Permalink
Add GoogleOpenAISchema with comprehensive test logging (#234)
Browse files Browse the repository at this point in the history
* Add GoogleOpenAISchema with OpenAI compatibility mode

- Add GoogleOpenAISchema struct with documentation
- Implement create_chat method with Google API base URL
- Implement create_embedding method with Google API base URL
- Use GOOGLE_API_KEY for authentication

* test: Add tests for GoogleOpenAISchema implementation

* test: Enhance GoogleOpenAISchema tests to improve coverage

* Move GoogleOpenAISchema tests to llm_openai_schema_def.jl

* feat: Add Gemini 1.5 models and aliases to user_preferences.jl

Added new models:
- gemini-1.5-pro-latest
- gemini-1.5-flash-8b-latest
- gemini-1.5-flash-latest

Added corresponding aliases:
- gem15p
- gem15f8
- gem15f

Set pricing according to Google's pay-as-you-go rates:
Input: /bin/bash.075/1M tokens
Output: /bin/bash.30/1M tokens

* Update Gemini model pricing to reflect correct per-million token rates

* revert: Remove unintended changes to test/llm_google.jl

* revert: restore test/llm_google.jl to original state from main branch

* feat: Add GoogleProvider with proper Bearer token auth and update GoogleOpenAISchema methods

- Add GoogleProvider struct with proper Bearer token authentication
- Update create_chat to use GoogleProvider and openai_request
- Update create_embeddings to use GoogleProvider and openai_request
- Maintain consistent URL handling up to /v1beta

* feat: Add Gemini models to registry and fix GoogleOpenAISchema tests

- Add Gemini 1.5 models (Pro, Flash, Flash 8b) with correct pricing
- Fix GoogleOpenAISchema tests to properly handle GOOGLE_API_KEY
- Save and restore original API key value during testing

* fix: restore gpt-4-turbo-preview model and ensure correct model ordering

* Add comprehensive logging to GoogleOpenAISchema test mock servers

- Add detailed request header logging
- Track authorization header values and expectations
- Log request body content and responses
- Improve debugging capabilities for test failures

* fix: Update auth_header in GoogleProvider to use OpenAIProvider implementation

* chore: prepare release v0.63.0

- Update version to 0.63.0
- Add warning about token/cost counting limitations in GoogleOpenAISchema
- Update changelog with correct model aliases (gem15p, gem15f, gem15f8)
- Remove logging from test/llm_openai_schema_def.jl

---------

Co-authored-by: devin-ai-integration[bot] <158243242+devin-ai-integration[bot]@users.noreply.github.com>
  • Loading branch information
devin-ai-integration[bot] authored Nov 12, 2024
1 parent 3cd535e commit faaf9ce
Show file tree
Hide file tree
Showing 8 changed files with 273 additions and 7 deletions.
6 changes: 6 additions & 0 deletions CHANGELOG.md
Original file line number Diff line number Diff line change
Expand Up @@ -10,6 +10,12 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0

### Fixed

## [0.63.0]

### Added
- Added support for Google's Gemini API via OpenAI compatibility mode (`GoogleOpenAISchema`). Use model aliases `gem15p` (Gemini 1.5 Pro), `gem15f` (Gemini 1.5 Flash), and `gem15f8` (Gemini 1.5 Flash 8b). Set your ENV api key `GOOGLE_API_KEY` to use it.
- Thanks to @sixzero, added support for Google Flash via OpenRouter and Qwen 72b models.

## [0.62.1]

### Fixed
Expand Down
2 changes: 1 addition & 1 deletion Project.toml
Original file line number Diff line number Diff line change
@@ -1,7 +1,7 @@
name = "PromptingTools"
uuid = "670122d1-24a8-4d70-bfce-740807c42192"
authors = ["J S @svilupp and contributors"]
version = "0.62.1"
version = "0.63.0"

[deps]
AbstractTrees = "1520ce14-60c1-5f80-bbc7-55ef81b5835c"
Expand Down
21 changes: 20 additions & 1 deletion src/llm_interface.jl
Original file line number Diff line number Diff line change
Expand Up @@ -274,6 +274,25 @@ Requires one environment variable to be set:
"""
struct XAIOpenAISchema <: AbstractOpenAISchema end

"""
GoogleOpenAISchema
Schema to call the Google's Gemini API using OpenAI compatibility mode. [API Reference](https://ai.google.dev/gemini-api/docs/openai#rest)
Links:
- [Get your API key](https://aistudio.google.com/apikey)
- [API Reference](https://ai.google.dev/gemini-api/docs/openai#rest)
- [Available models](https://ai.google.dev/models/gemini)
Requires one environment variable to be set:
- `GOOGLE_API_KEY`: Your API key
The base URL for the API is "https://generativelanguage.googleapis.com/v1beta"
Warning: Token counting and cost counting have not yet been implemented by Google, so you'll not have any such metrics. If you need it, use the native GoogleSchema with the GoogleGenAI.jl library.
"""
struct GoogleOpenAISchema <: AbstractOpenAISchema end

abstract type AbstractOllamaSchema <: AbstractPromptSchema end

"""
Expand Down Expand Up @@ -517,4 +536,4 @@ end
abstract type AbstractExtractedData end
Base.show(io::IO, x::AbstractExtractedData) = dump(io, x; maxdepth = 1)
"Check if the object is an instance of `AbstractExtractedData`"
isextracted(x) = x isa AbstractExtractedData
isextracted(x) = x isa AbstractExtractedData
45 changes: 44 additions & 1 deletion src/llm_openai_schema_defs.jl
Original file line number Diff line number Diff line change
Expand Up @@ -214,6 +214,34 @@ function OpenAI.create_chat(schema::XAIOpenAISchema,
api_key = isempty(XAI_API_KEY) ? api_key : XAI_API_KEY
OpenAI.create_chat(CustomOpenAISchema(), api_key, model, conversation; url, kwargs...)
end

# Add GoogleProvider implementation
Base.@kwdef struct GoogleProvider <: AbstractCustomProvider
api_key::String = ""
base_url::String = "https://generativelanguage.googleapis.com/v1beta"
api_version::String = ""
end

function OpenAI.auth_header(provider::GoogleProvider, api_key::AbstractString)
OpenAI.auth_header(OpenAI.OpenAIProvider(provider.api_key, provider.base_url, provider.api_version), api_key)
end

function OpenAI.create_chat(schema::GoogleOpenAISchema,
api_key::AbstractString,
model::AbstractString,
conversation;
url::String = "https://generativelanguage.googleapis.com/v1beta",
kwargs...)
api_key = isempty(GOOGLE_API_KEY) ? api_key : GOOGLE_API_KEY
# Use GoogleProvider instead of CustomProvider
provider = GoogleProvider(; api_key, base_url = url)
OpenAI.openai_request("chat/completions",
provider;
method = "POST",
messages = conversation,
model = model,
kwargs...)
end
function OpenAI.create_chat(schema::DatabricksOpenAISchema,
api_key::AbstractString,
model::AbstractString,
Expand Down Expand Up @@ -384,6 +412,21 @@ function OpenAI.create_embeddings(schema::XAIOpenAISchema,
base_url = url)
OpenAI.create_embeddings(provider, docs, model; kwargs...)
end
function OpenAI.create_embeddings(schema::GoogleOpenAISchema,
api_key::AbstractString,
docs,
model::AbstractString;
url::String = "https://generativelanguage.googleapis.com/v1beta",
kwargs...)
api_key = isempty(GOOGLE_API_KEY) ? api_key : GOOGLE_API_KEY
provider = GoogleProvider(; api_key, base_url = url)
OpenAI.openai_request("embeddings",
provider;
method = "POST",
input = docs,
model = model,
kwargs...)
end
function OpenAI.create_embeddings(schema::AzureOpenAISchema,
api_key::AbstractString,
docs,
Expand Down Expand Up @@ -441,4 +484,4 @@ function OpenAI.create_images(schema::TestEchoOpenAISchema,
schema.model_id = get(kwargs, :model, "")
schema.inputs = prompt
return schema
end
end
28 changes: 24 additions & 4 deletions src/user_preferences.jl
Original file line number Diff line number Diff line change
Expand Up @@ -474,7 +474,11 @@ aliases = merge(
"oro1" => "openai/o1-preview",
"oro1m" => "openai/o1-mini",
"orcop" => "cohere/command-r-plus-08-2024",
"orco" => "cohere/command-r-08-2024"
"orco" => "cohere/command-r-08-2024",
## Gemini 1.5 Models
"gem15p" => "gemini-1.5-pro-latest",
"gem15f8" => "gemini-1.5-flash-8b-latest",
"gem15f" => "gemini-1.5-flash-latest"
),
## Load aliases from preferences as well
@load_preference("MODEL_ALIASES", default=Dict{String, String}()))
Expand Down Expand Up @@ -509,12 +513,12 @@ registry = Dict{String, ModelSpec}(
OpenAISchema(),
1e-5,
3e-5,
"GPT-4 Turbo is an updated version of GPT4 that is much faster and the cheaper to use. 0125 refers to the release date of January 25, 2024."),
"GPT-4 Turbo is an updated version of GPT4 that is much faster and the cheaper to use. This is the general name for whatever is the latest GPT4 Turbo preview release. In April-24, it points to version 2024-04-09."),
"gpt-4-turbo" => ModelSpec("gpt-4-turbo",
OpenAISchema(),
1e-5,
3e-5,
"GPT-4 Turbo is an updated version of GPT4 that is much faster and the cheaper to use. This is the general name for whatever is the latest GPT4 Turbo preview release. In April-24, it points to version 2024-04-09."),
"GPT-4 Turbo is an updated version of GPT4 that is much faster and cheaper to use. This is the general name for whatever is the latest GPT4 Turbo preview release."),
"gpt-4-turbo-2024-04-09" => ModelSpec("gpt-4-turbo-2024-04-09",
OpenAISchema(),
1e-5,
Expand Down Expand Up @@ -1103,7 +1107,23 @@ registry = Dict{String, ModelSpec}(
XAIOpenAISchema(),
5e-6,
15e-6,
"XAI's Grok 2 beta model. Max 128K context.")
"XAI's Grok 2 beta model. Max 128K context."),
## Gemini 1.5 Models
"gemini-1.5-pro-latest" => ModelSpec("gemini-1.5-pro-latest",
GoogleOpenAISchema(),
1e-6,
5e-6,
"Gemini 1.5 Pro is Google's latest large language model with enhanced capabilities across reasoning, math, coding, and multilingual tasks. 128K context window."),
"gemini-1.5-flash-8b-latest" => ModelSpec("gemini-1.5-flash-8b-latest",
GoogleOpenAISchema(),
3.75e-8,
1.5e-7,
"Gemini 1.5 Flash 8B is a smaller, faster version of Gemini 1.5 optimized for quick responses while maintaining good performance. 128K context window."),
"gemini-1.5-flash-latest" => ModelSpec("gemini-1.5-flash-latest",
GoogleOpenAISchema(),
7.5e-8,
3.0e-7,
"Gemini 1.5 Flash is a high-performance model optimized for speed while maintaining strong capabilities across various tasks. 128K context window.")
)

"""
Expand Down
123 changes: 123 additions & 0 deletions test/llm_openai_schema_def.jl
Original file line number Diff line number Diff line change
@@ -0,0 +1,123 @@
using Test
using PromptingTools: GoogleOpenAISchema, AIMessage, aigenerate, aiembed

@testset "GoogleOpenAISchema" begin
# Save original API key
original_api_key = PromptingTools.GOOGLE_API_KEY


# Test with empty GOOGLE_API_KEY
PromptingTools.GOOGLE_API_KEY = ""
PORT = rand(10000:20000)
echo_server = HTTP.serve!(PORT, verbose = -1) do req
auth_header = HTTP.header(req, "Authorization")
@test HTTP.header(req, "Authorization") == "Bearer test_key"

content = JSON3.read(req.body)

response = Dict(
:choices => [
Dict(:message => Dict(:content => "Test response"),
:finish_reason => "stop")
],
:usage => Dict(:total_tokens => 5,
:prompt_tokens => 5,
:completion_tokens => 0))
return HTTP.Response(200, JSON3.write(response))
end

msg = aigenerate(GoogleOpenAISchema(),
"Test prompt";
api_key = "test_key",
model = "gemini-1.5-pro-latest",
api_kwargs = (; url = "http://localhost:$(PORT)"))

@test msg.content == "Test response"
@test msg.finish_reason == "stop"
close(echo_server)

# Test with non-empty GOOGLE_API_KEY
PromptingTools.GOOGLE_API_KEY = "env_key"
PORT = rand(10000:20000)
echo_server = HTTP.serve!(PORT, verbose = -1) do req
auth_header = HTTP.header(req, "Authorization")
@test HTTP.header(req, "Authorization") == "Bearer env_key"

content = JSON3.read(req.body)

response = Dict(
:choices => [
Dict(:message => Dict(:content => "Test response"),
:finish_reason => "stop")
],
:usage => Dict(:total_tokens => 5,
:prompt_tokens => 5,
:completion_tokens => 0))
return HTTP.Response(200, JSON3.write(response))
end

msg = aigenerate(GoogleOpenAISchema(),
"Test prompt";
api_key = "test_key", # This should be ignored since GOOGLE_API_KEY is set
model = "gemini-1.5-pro-latest",
api_kwargs = (; url = "http://localhost:$(PORT)"))

@test msg.content == "Test response"
@test msg.finish_reason == "stop"
close(echo_server)

# Test embeddings with empty GOOGLE_API_KEY
PromptingTools.GOOGLE_API_KEY = ""
PORT = rand(10000:20000)
echo_server = HTTP.serve!(PORT, verbose = -1) do req
auth_header = HTTP.header(req, "Authorization")
@test HTTP.header(req, "Authorization") == "Bearer test_key"

content = JSON3.read(req.body)

response = Dict(:data => [Dict(:embedding => ones(128))],
:usage => Dict(:total_tokens => 5,
:prompt_tokens => 5,
:completion_tokens => 0))
return HTTP.Response(200, JSON3.write(response))
end

msg = aiembed(GoogleOpenAISchema(),
"Test prompt";
api_key = "test_key",
model = "gemini-1.5-pro-latest",
api_kwargs = (; url = "http://localhost:$(PORT)"))

@test msg.content == ones(128)
@test msg.tokens == (5, 0)
close(echo_server)

# Test embeddings with non-empty GOOGLE_API_KEY
PromptingTools.GOOGLE_API_KEY = "env_key"
PORT = rand(10000:20000)
echo_server = HTTP.serve!(PORT, verbose = -1) do req
auth_header = HTTP.header(req, "Authorization")
@test HTTP.header(req, "Authorization") == "Bearer env_key"

content = JSON3.read(req.body)

response = Dict(:data => [Dict(:embedding => ones(128))],
:usage => Dict(:total_tokens => 5,
:prompt_tokens => 5,
:completion_tokens => 0))
return HTTP.Response(200, JSON3.write(response))
end

msg = aiembed(GoogleOpenAISchema(),
"Test prompt";
api_key = "test_key", # This should be ignored since GOOGLE_API_KEY is set
model = "gemini-1.5-pro-latest",
api_kwargs = (; url = "http://localhost:$(PORT)"))

@test msg.content == ones(128)
@test msg.tokens == (5, 0)
close(echo_server)

# Restore original API key
PromptingTools.GOOGLE_API_KEY = original_api_key
end
1 change: 1 addition & 0 deletions test/runtests.jl
Original file line number Diff line number Diff line change
Expand Up @@ -28,6 +28,7 @@ end
include("llm_ollama_managed.jl")
include("llm_ollama.jl")
include("llm_google.jl")
include("llm_openai_schema_def.jl")
include("llm_anthropic.jl")
include("llm_sharegpt.jl")
include("llm_tracer.jl")
Expand Down
54 changes: 54 additions & 0 deletions test_analysis.txt
Original file line number Diff line number Diff line change
@@ -0,0 +1,54 @@
Test Coverage Analysis for GoogleOpenAISchema

Current Issues:
- 6 uncovered lines in src/llm_openai_schema_defs.jl
- Patch coverage at 14.29%

Implementation Analysis:
1. GoogleProvider struct (lines 219-223)
- Basic struct definition, likely covered
- auth_header method (lines 225-227) might be uncovered

2. create_chat method (lines 229-244)
- Key lines that might be uncovered:
* Line 235: api_key override logic
* Line 237: GoogleProvider instantiation
* Lines 238-243: OpenAI.openai_request call

3. create_embeddings method (lines 415-429)
- Similar pattern to create_chat
- Potential uncovered lines:
* Line 421: api_key override logic
* Line 422: GoogleProvider instantiation
* Lines 423-428: OpenAI.openai_request call

Hypotheses for Coverage Issues:
1. Streaming callback paths not tested
- We're using openai_request directly, which might have different behavior
- Solution: Add tests for streaming scenarios

2. Error handling paths not tested
- No tests for API errors or invalid responses
- Solution: Add tests with mock error responses

3. Provider instantiation edge cases
- GoogleProvider creation with different URL combinations not tested
- Solution: Add tests with various URL configurations

4. API key override logic not fully tested
- Need to test all combinations of empty/non-empty GOOGLE_API_KEY
- Solution: Expand current tests to cover more scenarios

5. Request parameter handling not fully tested
- Different combinations of optional parameters not tested
- Solution: Add tests with various kwargs combinations

Most Likely Issue:
Hypothesis #4 seems most likely - our tests don't fully exercise the API key override logic in both create_chat and create_embeddings methods. The current tests only check basic scenarios, but we need to test edge cases and different combinations of API keys.

Action Plan:
1. Add tests for API key override edge cases
2. Add tests for different URL configurations
3. Add tests for error scenarios
4. Add tests for streaming callbacks
5. Add tests for various kwargs combinations

0 comments on commit faaf9ce

Please sign in to comment.