This commit is contained in:
+75
-2
@@ -1,5 +1,7 @@
|
||||
use reqwest::Client;
|
||||
use serde_json::Value;
|
||||
use serde_json::{Value, json};
|
||||
|
||||
use crate::errors::OllamaError;
|
||||
|
||||
#[derive(Clone)]
|
||||
pub struct OllamaProvider {
|
||||
@@ -15,11 +17,82 @@ impl OllamaProvider {
|
||||
}
|
||||
}
|
||||
|
||||
pub async fn list_models(&self) -> Result<Value, reqwest::Error> {
|
||||
pub async fn list_models(&self) -> Result<Value, OllamaError> {
|
||||
let url = format!("{}/api/tags", self.base_url);
|
||||
|
||||
let res = self.client.get(url).send().await?.json::<Value>().await?;
|
||||
|
||||
Ok(res)
|
||||
}
|
||||
|
||||
pub async fn completions(&self, body: Value) -> Result<Value, OllamaError> {
|
||||
// Validate prompt
|
||||
let prompt = body
|
||||
.get("prompt")
|
||||
.and_then(|v| v.as_str())
|
||||
.filter(|s| !s.trim().is_empty())
|
||||
.ok_or(OllamaError::MissingPrompt)?;
|
||||
|
||||
let model = body
|
||||
.get("model")
|
||||
.and_then(|v| v.as_str())
|
||||
.unwrap_or("llama3");
|
||||
|
||||
// Validate model availability
|
||||
let available = self.list_models().await?; // Http error auto-converts via #[from]
|
||||
let model_exists = available
|
||||
.get("models")
|
||||
.and_then(|m| m.as_array())
|
||||
.map(|arr| {
|
||||
arr.iter()
|
||||
.any(|m| m.get("name").and_then(|n| n.as_str()) == Some(model))
|
||||
})
|
||||
.unwrap_or(false);
|
||||
|
||||
if !model_exists {
|
||||
return Err(OllamaError::ModelNotFound(model.to_string()));
|
||||
}
|
||||
|
||||
// Build payload (now using validated locals)
|
||||
let ollama_payload = json!({
|
||||
"model": model,
|
||||
"prompt": prompt,
|
||||
"stream": false,
|
||||
"options": {
|
||||
"temperature": body.get("temperature"),
|
||||
"top_p": body.get("top_p"),
|
||||
"num_predict": body.get("max_tokens"),
|
||||
}
|
||||
});
|
||||
|
||||
let url = format!("{}/api/generate", self.base_url);
|
||||
let res = self
|
||||
.client
|
||||
.post(&url)
|
||||
.json(&ollama_payload)
|
||||
.send()
|
||||
.await?
|
||||
.json::<Value>()
|
||||
.await?;
|
||||
|
||||
Ok(json!({
|
||||
"id": "cmpl-ollama",
|
||||
"object": "text_completion",
|
||||
"model": res.get("model"),
|
||||
"choices": [{
|
||||
"text": res.get("response"),
|
||||
"index": 0,
|
||||
"finish_reason": if res.get("done").and_then(|v| v.as_bool()).unwrap_or(false) {
|
||||
"stop"
|
||||
} else {
|
||||
"length"
|
||||
},
|
||||
}],
|
||||
"usage": {
|
||||
"prompt_tokens": res.get("prompt_eval_count"),
|
||||
"completion_tokens": res.get("eval_count"),
|
||||
"total_tokens": null,
|
||||
}
|
||||
}))
|
||||
}
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user