Skip to content

Commit 5304adc

Browse files
google-labs-jules[bot]Tcode-Motion
google-labs-jules[bot]
and
Tcode-Motion
committed
Extract AI provider logic into helper functions to reduce nesting in register_ai
Co-authored-by: Tcode-Motion <188012755+Tcode-Motion@users.noreply.github.com>
1 parent 514e6c0 commit 5304adc

1 file changed

Lines changed: 149 additions & 79 deletions

File tree

‎stdlib/src/ai.rs‎

Lines changed: 149 additions & 79 deletions
Original file line numberDiff line numberDiff line change
@@ -26,7 +26,10 @@ impl StdlibRegistry {
2626
// Check environment capability for retrieving API keys
2727
if !ctx.config.capabilities.contains(&Capability::Environment) {
2828
return Err(RuntimeError::new(
29-
RuntimeErrorKind::InvalidOperation("Security policy violation: Environment capability is denied".to_string()),
29+
RuntimeErrorKind::InvalidOperation(
30+
"Security policy violation: Environment capability is denied"
31+
.to_string(),
32+
),
3033
None,
3134
None,
3235
));
@@ -35,93 +38,27 @@ impl StdlibRegistry {
3538
// Also requires Network capability to make HTTP requests
3639
if !ctx.config.capabilities.contains(&Capability::Network) {
3740
return Err(RuntimeError::new(
38-
RuntimeErrorKind::InvalidOperation("Security policy violation: Network capability is denied".to_string()),
41+
RuntimeErrorKind::InvalidOperation(
42+
"Security policy violation: Network capability is denied"
43+
.to_string(),
44+
),
3945
None,
4046
None,
4147
));
4248
}
4349

4450
match provider.as_str() {
45-
"openai" => {
46-
let key = std::env::var("OPENAI_API_KEY").unwrap_or_default();
47-
if key.is_empty() {
48-
return Ok(RuntimeValue::Str(format!("[Mock OpenAI Response] Prompt: {}", prompt)));
49-
}
50-
51-
// Real HTTP call to OpenAI Chat Completion
52-
let body = serde_json::json!({
53-
"model": "gpt-4o-mini",
54-
"messages": [{"role": "user", "content": prompt}]
55-
});
56-
57-
let resp = ureq::post("https://api.openai.com/v1/chat/completions")
58-
.set("Authorization", &format!("Bearer {}", key))
59-
.set("Content-Type", "application/json")
60-
.send_json(body)
61-
.map_err(|e| RuntimeError::new(RuntimeErrorKind::InvalidOperation(format!("OpenAI request failed: {}", e)), None, None))?;
62-
63-
let json: serde_json::Value = resp.into_json()
64-
.map_err(|e| RuntimeError::new(RuntimeErrorKind::InvalidOperation(format!("Failed to parse OpenAI JSON response: {}", e)), None, None))?;
65-
66-
let content = json["choices"][0]["message"]["content"].as_str()
67-
.ok_or_else(|| RuntimeError::new(RuntimeErrorKind::InvalidOperation("OpenAI response content empty".to_string()), None, None))?;
68-
69-
Ok(RuntimeValue::Str(content.to_string()))
70-
}
71-
"gemini" => {
72-
let key = std::env::var("GEMINI_API_KEY").unwrap_or_default();
73-
if key.is_empty() {
74-
return Ok(RuntimeValue::Str(format!("[Mock Gemini Response] Prompt: {}", prompt)));
75-
}
76-
77-
// Real HTTP call to Gemini API
78-
let url = format!("https://generativelanguage.googleapis.com/v1beta/models/gemini-1.5-flash:generateContent?key={}", key);
79-
let body = serde_json::json!({
80-
"contents": [{
81-
"parts": [{"text": prompt}]
82-
}]
83-
});
84-
85-
let resp = ureq::post(&url)
86-
.set("Content-Type", "application/json")
87-
.send_json(body)
88-
.map_err(|e| RuntimeError::new(RuntimeErrorKind::InvalidOperation(format!("Gemini request failed: {}", e)), None, None))?;
89-
90-
let json: serde_json::Value = resp.into_json()
91-
.map_err(|e| RuntimeError::new(RuntimeErrorKind::InvalidOperation(format!("Failed to parse Gemini JSON response: {}", e)), None, None))?;
92-
93-
let content = json["candidates"][0]["content"]["parts"][0]["text"].as_str()
94-
.ok_or_else(|| RuntimeError::new(RuntimeErrorKind::InvalidOperation("Gemini response content empty".to_string()), None, None))?;
95-
96-
Ok(RuntimeValue::Str(content.to_string()))
97-
}
98-
"local" => {
99-
// Mock local Llama.cpp inference endpoint check (e.g. running on localhost:8080)
100-
let local_url = "http://127.0.0.1:8080/completion";
101-
let body = serde_json::json!({
102-
"prompt": prompt,
103-
"n_predict": 128
104-
});
105-
106-
match ureq::post(local_url).set("Content-Type", "application/json").send_json(body) {
107-
Ok(resp) => {
108-
if let Ok(json) = resp.into_json::<serde_json::Value>() {
109-
if let Some(content) = json["content"].as_str() {
110-
return Ok(RuntimeValue::Str(content.to_string()));
111-
}
112-
}
113-
Ok(RuntimeValue::Str("[Mock Local LLM Response] (local server responded with invalid content)".to_string()))
114-
}
115-
Err(_) => {
116-
Ok(RuntimeValue::Str(format!("[Mock Local LLM Response] Prompt: {}", prompt)))
117-
}
118-
}
119-
}
51+
"openai" => generate_openai(&prompt),
52+
"gemini" => generate_gemini(&prompt),
53+
"local" => generate_local(&prompt),
12054
_ => Err(RuntimeError::new(
121-
RuntimeErrorKind::InvalidOperation(format!("Unknown AI provider: {}", provider)),
55+
RuntimeErrorKind::InvalidOperation(format!(
56+
"Unknown AI provider: {}",
57+
provider
58+
)),
12259
None,
12360
None,
124-
))
61+
)),
12562
}
12663
},
12764
}),
@@ -138,3 +75,136 @@ impl StdlibRegistry {
13875
);
13976
}
14077
}
78+
79+
fn generate_openai(prompt: &str) -> Result<RuntimeValue, RuntimeError> {
80+
let key = std::env::var("OPENAI_API_KEY").unwrap_or_default();
81+
if key.is_empty() {
82+
return Ok(RuntimeValue::Str(format!(
83+
"[Mock OpenAI Response] Prompt: {}",
84+
prompt
85+
)));
86+
}
87+
88+
// Real HTTP call to OpenAI Chat Completion
89+
let body = serde_json::json!({
90+
"model": "gpt-4o-mini",
91+
"messages": [{"role": "user", "content": prompt}]
92+
});
93+
94+
let resp = ureq::post("https://api.openai.com/v1/chat/completions")
95+
.set("Authorization", &format!("Bearer {}", key))
96+
.set("Content-Type", "application/json")
97+
.send_json(body)
98+
.map_err(|e| {
99+
RuntimeError::new(
100+
RuntimeErrorKind::InvalidOperation(format!("OpenAI request failed: {}", e)),
101+
None,
102+
None,
103+
)
104+
})?;
105+
106+
let json: serde_json::Value = resp.into_json().map_err(|e| {
107+
RuntimeError::new(
108+
RuntimeErrorKind::InvalidOperation(format!(
109+
"Failed to parse OpenAI JSON response: {}",
110+
e
111+
)),
112+
None,
113+
None,
114+
)
115+
})?;
116+
117+
let content = json["choices"][0]["message"]["content"]
118+
.as_str()
119+
.ok_or_else(|| {
120+
RuntimeError::new(
121+
RuntimeErrorKind::InvalidOperation("OpenAI response content empty".to_string()),
122+
None,
123+
None,
124+
)
125+
})?;
126+
127+
Ok(RuntimeValue::Str(content.to_string()))
128+
}
129+
130+
fn generate_gemini(prompt: &str) -> Result<RuntimeValue, RuntimeError> {
131+
let key = std::env::var("GEMINI_API_KEY").unwrap_or_default();
132+
if key.is_empty() {
133+
return Ok(RuntimeValue::Str(format!(
134+
"[Mock Gemini Response] Prompt: {}",
135+
prompt
136+
)));
137+
}
138+
139+
// Real HTTP call to Gemini API
140+
let url = format!("https://generativelanguage.googleapis.com/v1beta/models/gemini-1.5-flash:generateContent?key={}", key);
141+
let body = serde_json::json!({
142+
"contents": [{
143+
"parts": [{"text": prompt}]
144+
}]
145+
});
146+
147+
let resp = ureq::post(&url)
148+
.set("Content-Type", "application/json")
149+
.send_json(body)
150+
.map_err(|e| {
151+
RuntimeError::new(
152+
RuntimeErrorKind::InvalidOperation(format!("Gemini request failed: {}", e)),
153+
None,
154+
None,
155+
)
156+
})?;
157+
158+
let json: serde_json::Value = resp.into_json().map_err(|e| {
159+
RuntimeError::new(
160+
RuntimeErrorKind::InvalidOperation(format!(
161+
"Failed to parse Gemini JSON response: {}",
162+
e
163+
)),
164+
None,
165+
None,
166+
)
167+
})?;
168+
169+
let content = json["candidates"][0]["content"]["parts"][0]["text"]
170+
.as_str()
171+
.ok_or_else(|| {
172+
RuntimeError::new(
173+
RuntimeErrorKind::InvalidOperation("Gemini response content empty".to_string()),
174+
None,
175+
None,
176+
)
177+
})?;
178+
179+
Ok(RuntimeValue::Str(content.to_string()))
180+
}
181+
182+
fn generate_local(prompt: &str) -> Result<RuntimeValue, RuntimeError> {
183+
// Mock local Llama.cpp inference endpoint check (e.g. running on localhost:8080)
184+
let local_url = "http://127.0.0.1:8080/completion";
185+
let body = serde_json::json!({
186+
"prompt": prompt,
187+
"n_predict": 128
188+
});
189+
190+
match ureq::post(local_url)
191+
.set("Content-Type", "application/json")
192+
.send_json(body)
193+
{
194+
Ok(resp) => {
195+
if let Ok(json) = resp.into_json::<serde_json::Value>() {
196+
if let Some(content) = json["content"].as_str() {
197+
return Ok(RuntimeValue::Str(content.to_string()));
198+
}
199+
}
200+
Ok(RuntimeValue::Str(
201+
"[Mock Local LLM Response] (local server responded with invalid content)"
202+
.to_string(),
203+
))
204+
}
205+
Err(_) => Ok(RuntimeValue::Str(format!(
206+
"[Mock Local LLM Response] Prompt: {}",
207+
prompt
208+
))),
209+
}
210+
}

0 commit comments

Comments
 (0)