rusty_commit/providers/
nvidia.rs

1//! NVIDIA NIM Provider - Enterprise GPU Inference
2//!
3//! NVIDIA NIM (NVIDIA Inference Microservices) provides optimized inference
4//! for LLMs on NVIDIA GPUs. Supports both self-hosted and cloud deployments.
5//!
6//! Setup:
7//! 1. Get API key from: https://build.nvidia.com
8//! 2. Configure rco:
9//!    `rco config set RCO_AI_PROVIDER=nvidia RCO_API_KEY=<key> RCO_MODEL=meta/llama-3.1-8b-instruct`
10//!
11//! Docs: https://docs.nvidia.com/nim/
12
13use anyhow::{Context, Result};
14use async_trait::async_trait;
15use reqwest::Client;
16use serde::{Deserialize, Serialize};
17
18use super::{build_prompt, AIProvider};
19use crate::config::Config;
20use crate::utils::retry::retry_async;
21
22pub struct NvidiaProvider {
23    client: Client,
24    api_url: String,
25    api_key: String,
26    model: String,
27}
28
29#[derive(Serialize)]
30struct NvidiaRequest {
31    model: String,
32    messages: Vec<NvidiaMessage>,
33    max_tokens: i32,
34    temperature: f32,
35    top_p: f32,
36    stream: bool,
37}
38
39#[derive(Serialize, Deserialize, Clone)]
40struct NvidiaMessage {
41    role: String,
42    content: String,
43}
44
45#[derive(Deserialize)]
46struct NvidiaResponse {
47    choices: Vec<NvidiaChoice>,
48}
49
50#[derive(Deserialize)]
51struct NvidiaChoice {
52    message: NvidiaMessage,
53}
54
55impl NvidiaProvider {
56    pub fn new(config: &Config) -> Result<Self> {
57        let client = Client::new();
58        let api_key = config
59            .api_key
60            .as_ref()
61            .context("NVIDIA API key not configured.\nRun: rco config set RCO_API_KEY=<your_key>\nGet your API key from: https://build.nvidia.com")?;
62
63        let api_url = config
64            .api_url
65            .as_deref()
66            .unwrap_or("https://integrate.api.nvidia.com/v1")
67            .to_string();
68
69        let model = config
70            .model
71            .as_deref()
72            .unwrap_or("meta/llama-3.1-8b-instruct")
73            .to_string();
74
75        Ok(Self {
76            client,
77            api_url,
78            api_key: api_key.clone(),
79            model,
80        })
81    }
82
83    /// Create provider from account configuration
84    #[allow(dead_code)]
85    pub fn from_account(
86        account: &crate::config::accounts::AccountConfig,
87        api_key: &str,
88        config: &Config,
89    ) -> Result<Self> {
90        let client = Client::new();
91        let api_url = account
92            .api_url
93            .as_deref()
94            .or(config.api_url.as_deref())
95            .unwrap_or("https://integrate.api.nvidia.com/v1")
96            .to_string();
97
98        let model = account
99            .model
100            .as_deref()
101            .or(config.model.as_deref())
102            .unwrap_or("meta/llama-3.1-8b-instruct")
103            .to_string();
104
105        Ok(Self {
106            client,
107            api_url,
108            api_key: api_key.to_string(),
109            model,
110        })
111    }
112}
113
114#[async_trait]
115impl AIProvider for NvidiaProvider {
116    async fn generate_commit_message(
117        &self,
118        diff: &str,
119        context: Option<&str>,
120        full_gitmoji: bool,
121        config: &Config,
122    ) -> Result<String> {
123        let prompt = build_prompt(diff, context, config, full_gitmoji);
124
125        let messages = vec![
126            NvidiaMessage {
127                role: "system".to_string(),
128                content: "You are an expert at writing clear, concise git commit messages.".to_string(),
129            },
130            NvidiaMessage {
131                role: "user".to_string(),
132                content: prompt,
133            },
134        ];
135
136        let request = NvidiaRequest {
137            model: self.model.clone(),
138            messages,
139            max_tokens: config.tokens_max_output.unwrap_or(500) as i32,
140            temperature: 0.7,
141            top_p: 0.7,
142            stream: false,
143        };
144
145        let nvidia_response: NvidiaResponse = retry_async(|| async {
146            let url = format!("{}/chat/completions", self.api_url);
147            let response = self
148                .client
149                .post(&url)
150                .header("Authorization", format!("Bearer {}", self.api_key))
151                .json(&request)
152                .send()
153                .await
154                .context("Failed to connect to NVIDIA NIM API")?;
155
156            if !response.status().is_success() {
157                let error_text = response.text().await?;
158                if error_text.contains("401") || error_text.contains("Unauthorized") {
159                    return Err(anyhow::anyhow!("Invalid NVIDIA API key. Please check your API key configuration."));
160                }
161                return Err(anyhow::anyhow!("NVIDIA NIM API error: {}", error_text));
162            }
163
164            let nvidia_response: NvidiaResponse = response
165                .json()
166                .await
167                .context("Failed to parse NVIDIA NIM response")?;
168
169            Ok(nvidia_response)
170        })
171        .await
172        .context("Failed to generate commit message from NVIDIA NIM after retries")?;
173
174        let message = nvidia_response
175            .choices
176            .first()
177            .map(|choice| choice.message.content.trim().to_string())
178            .context("NVIDIA NIM returned an empty response")?;
179
180        Ok(message)
181    }
182}
183
184/// ProviderBuilder for NVIDIA NIM
185pub struct NvidiaProviderBuilder;
186
187impl super::registry::ProviderBuilder for NvidiaProviderBuilder {
188    fn name(&self) -> &'static str {
189        "nvidia"
190    }
191
192    fn aliases(&self) -> Vec<&'static str> {
193        vec!["nvidia-nim", "nim", "nvidia-ai"]
194    }
195
196    fn category(&self) -> super::registry::ProviderCategory {
197        super::registry::ProviderCategory::Cloud
198    }
199
200    fn create(&self, config: &Config) -> Result<Box<dyn super::AIProvider>> {
201        Ok(Box::new(NvidiaProvider::new(config)?))
202    }
203
204    fn requires_api_key(&self) -> bool {
205        true
206    }
207
208    fn default_model(&self) -> Option<&'static str> {
209        Some("meta/llama-3.1-8b-instruct")
210    }
211}
rusty_commit/providers/nvidia.rs

rusty_commit/providers/
nvidia.rs
