Skip to content

Configure models and API base from environment #2

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 3 commits into from
Apr 27, 2025
Merged

Configure models and API base from environment #2

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
cf8b570
feat: get api base from env
pashadia Apr 18, 2025
aafff22
feat: get embedding model from env
pashadia Apr 18, 2025
e2294df
feat: get llm model from env
pashadia Apr 18, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
19 changes: 11 additions & 8 deletions src/main.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -11,7 +11,7 @@ use crate::{
error::ServerError,
server::RustDocsServer, // Import the updated RustDocsServer
};
use async_openai::Client as OpenAIClient;
use async_openai::{Client as OpenAIClient, config::OpenAIConfig};
use bincode::config;
use cargo::core::PackageIdSpec;
use clap::Parser; // Import clap Parser
Expand Down Expand Up @@ -182,7 +182,12 @@ async fn main() -> Result<(), ServerError> {
let mut documents_for_server: Vec<Document> = loaded_documents_from_cache.unwrap_or_default();

// --- Initialize OpenAI Client (needed for question embedding even if cache hit) ---
let openai_client = OpenAIClient::new();
let openai_client = if let Ok(api_base) = env::var("OPENAI_API_BASE") {
let config = OpenAIConfig::new().with_api_base(api_base);
OpenAIClient::with_config(config)
} else {
OpenAIClient::new()
};
OPENAI_CLIENT
.set(openai_client.clone()) // Clone the client for the OnceCell
.expect("Failed to set OpenAI client");
Expand All @@ -209,12 +214,10 @@ async fn main() -> Result<(), ServerError> {
documents_for_server = loaded_documents.clone();

eprintln!("Generating embeddings...");
let (generated_embeddings, total_tokens) = generate_embeddings(
&openai_client,
&loaded_documents,
"text-embedding-3-small",
)
.await?;
let embedding_model: String = env::var("EMBEDDING_MODEL")
.unwrap_or_else(|_| "text-embedding-3-small".to_string());
let (generated_embeddings, total_tokens) =
generate_embeddings(&openai_client, &loaded_documents, &embedding_model).await?;

let cost_per_million = 0.02;
let estimated_cost = (total_tokens as f64 / 1_000_000.0) * cost_per_million;
Expand Down
9 changes: 6 additions & 3 deletions src/server.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -175,8 +175,10 @@ impl RustDocsServer {
.get()
.ok_or_else(|| McpError::internal_error("OpenAI client not initialized", None))?;

let embedding_model: String =
env::var("EMBEDDING_MODEL").unwrap_or_else(|_| "text-embedding-3-small".to_string());
let question_embedding_request = CreateEmbeddingRequestArgs::default()
.model("text-embedding-3-small")
.model(embedding_model)
.input(question.to_string())
.build()
.map_err(|e| {
Expand Down Expand Up @@ -223,8 +225,10 @@ impl RustDocsServer {
doc.content, question
);

let llm_model: String = env::var("LLM_MODEL")
.unwrap_or_else(|_| "gpt-4o-mini-2024-07-18".to_string());
let chat_request = CreateChatCompletionRequestArgs::default()
.model("gpt-4o-mini-2024-07-18")
.model(llm_model)
.messages(vec![
ChatCompletionRequestSystemMessageArgs::default()
.content(system_prompt)
Expand Down Expand Up @@ -379,5 +383,4 @@ impl ServerHandler for RustDocsServer {
resource_templates: Vec::new(), // No templates defined yet
})
}

}