Skip to content
This repository has been archived by the owner on Aug 2, 2024. It is now read-only.

Commit

Permalink
Refactor OpenAI config and prompt functions.
Browse files Browse the repository at this point in the history
- Use OpenAI config in `Client` initialization
- Refactor `with_api_key()` method to `OpenAIConfig::with_api_key()`
- Add generic type parameter to `chat_completion` and `completion` functions in `prompt.rs`
  • Loading branch information
zurawiki committed Jun 17, 2023
1 parent d5f1ca1 commit 4407711
Show file tree
Hide file tree
Showing 2 changed files with 8 additions and 5 deletions.
4 changes: 3 additions & 1 deletion src/cli.rs
Original file line number Diff line number Diff line change
@@ -1,5 +1,6 @@
use crate::prompt;
use anyhow::bail;
use async_openai::config::OpenAIConfig;
use async_openai::Client;
use clap::{command, value_parser, Parser};
use std::fs::File;
Expand Down Expand Up @@ -69,7 +70,8 @@ pub(crate) async fn main() -> anyhow::Result<()> {
bail!("OPENAI_API_KEY must be set");
};

let client = Client::new().with_backoff(backoff).with_api_key(api_key);
let client = Client::<OpenAIConfig>::with_config(OpenAIConfig::new().with_api_key(api_key))
.with_backoff(backoff);

let input = get_prompt_from_input(&cli.files)?;

Expand Down
9 changes: 5 additions & 4 deletions src/prompt.rs
Original file line number Diff line number Diff line change
@@ -1,5 +1,6 @@
use anyhow::bail;
use async_openai::{
config::Config,
types::{
ChatCompletionRequestMessageArgs, CreateChatCompletionRequestArgs,
CreateCompletionRequestArgs, Role,
Expand All @@ -16,8 +17,8 @@ pub(crate) fn should_use_chat_completion(model: &str) -> bool {
model.to_lowercase().starts_with("gpt-4") || model.to_lowercase().starts_with("gpt-3.5-turbo")
}

pub(crate) async fn chat_completion(
client: &Client,
pub(crate) async fn chat_completion<C: Config>(
client: &Client<C>,
prompt: &str,
model: &str,
cli: &CompletionArgs,
Expand Down Expand Up @@ -74,8 +75,8 @@ pub(crate) async fn chat_completion(

Ok(())
}
pub(crate) async fn completion(
client: &Client,
pub(crate) async fn completion<C: Config>(
client: &Client<C>,
prompt: &str,
model: &str,
cli: &CompletionArgs,
Expand Down

0 comments on commit 4407711

Please sign in to comment.