From fe4f4b67aa2c98b6409ed2d96df23dca1e2c8530 Mon Sep 17 00:00:00 2001
From: lux <aero.lu@cyte.site>
Date: Thu, 13 Jun 2024 03:58:06 +0800
Subject: [PATCH] suggested changes

---
 src/config.rs                      |  6 +++---
 src/transformer_backends/ollama.rs | 21 ++++++++++-----------
 2 files changed, 13 insertions(+), 14 deletions(-)

diff --git a/src/config.rs b/src/config.rs
index 33e35b8..82df19e 100644
--- a/src/config.rs
+++ b/src/config.rs
@@ -94,9 +94,9 @@ const fn n_ctx_default() -> u32 {
 #[derive(Clone, Debug, Deserialize)]
 #[serde(deny_unknown_fields)]
 pub struct Ollama {
-    // The completions endpoint
-    pub completions_endpoint: Option<String>,
-    // The chat endpoint
+    // The generate endpoint, default: 'http://localhost:11434/api/generate'
+    pub generate_endpoint: Option<String>,
+    // The chat endpoint, default: 'http://localhost:11434/api/chat'
     pub chat_endpoint: Option<String>,
     // The model name
     pub model: String,
diff --git a/src/transformer_backends/ollama.rs b/src/transformer_backends/ollama.rs
index 7d86f0f..c1de053 100644
--- a/src/transformer_backends/ollama.rs
+++ b/src/transformer_backends/ollama.rs
@@ -1,4 +1,3 @@
-use anyhow::Context;
 use serde::{Deserialize, Serialize};
 use serde_json::{json, Value};
 use std::collections::HashMap;
@@ -68,11 +67,11 @@ impl Ollama {
     ) -> anyhow::Result<String> {
         let client = reqwest::Client::new();
         let res: OllamaCompletionsResponse = client
-            .post(
-                self.configuration
-                    .completions_endpoint
-                    .as_ref()
-                    .context("specify `completions_endpoint` to use completions. Wanted to use `chat` instead? Please specify `chat_endpoint` and `messages`.")?,
+            .post(self
+                .configuration
+                .generate_endpoint
+                .as_ref()
+                .unwrap_or(&"http://localhost:11434/api/generate".to_string())
             )
             .header("Content-Type", "application/json")
             .header("Accept", "application/json")
@@ -107,11 +106,11 @@ impl Ollama {
     ) -> anyhow::Result<String> {
         let client = reqwest::Client::new();
         let res: OllamaChatResponse = client
-            .post(
-                self.configuration
-                    .chat_endpoint
-                    .as_ref()
-                    .context("must specify `chat_endpoint` to use chat")?,
+            .post(self
+                .configuration
+                .chat_endpoint
+                .as_ref()
+                .unwrap_or(&"http://localhost:11434/api/chat".to_string())
             )
             .header("Content-Type", "application/json")
             .header("Accept", "application/json")