From 176ff6f83ecfcf044a7deffad61a33b43480cdc5 Mon Sep 17 00:00:00 2001
From: Kelvie Wong <kelvie@kelvie.ca>
Date: Tue, 30 Jan 2024 03:43:55 -0800
Subject: feat: add `extra_fields` to models of localai/ollama clients (#298)

* Add an "extra_fields" config to localai models

Because there are so many local AIs out there with a bunch of custom
parameters you can set, this allows users to send in extra parameters to
a local LLM runner, such as, e.g. `instruction_template: Alpaca`, so
that Mixtral can take a system prompt.

* support ollama

---------

Co-authored-by: sigoden <sigoden@gmail.com>
---
 src/client/model.rs | 21 +++++++++++++++++++++
 1 file changed, 21 insertions(+)

(limited to 'src/client/model.rs')

diff --git a/src/client/model.rs b/src/client/model.rs
index 88d3294..b29166e 100644
--- a/src/client/model.rs
+++ b/src/client/model.rs
@@ -12,6 +12,7 @@ pub struct Model {
     pub client_name: String,
     pub name: String,
     pub max_tokens: Option<usize>,
+    pub extra_fields: Option<serde_json::Map<String, serde_json::Value>>,
     pub tokens_count_factors: TokensCountFactors,
     pub capabilities: ModelCapabilities,
 }
@@ -27,6 +28,7 @@ impl Model {
         Self {
             client_name: client_name.into(),
             name: name.into(),
+            extra_fields: None,
             max_tokens: None,
             tokens_count_factors: Default::default(),
             capabilities: ModelCapabilities::Text,
@@ -73,6 +75,14 @@ impl Model {
         self
     }
 
+    pub fn set_extra_fields(
+        mut self,
+        extra_fields: Option<serde_json::Map<String, serde_json::Value>>,
+    ) -> Self {
+        self.extra_fields = extra_fields;
+        self
+    }
+
     pub fn set_max_tokens(mut self, max_tokens: Option<usize>) -> Self {
         match max_tokens {
             None | Some(0) => self.max_tokens = None,
@@ -122,12 +132,23 @@ impl Model {
         }
         Ok(())
     }
+
+    pub fn merge_extra_fields(&self, body: &mut serde_json::Value) {
+        if let (Some(body), Some(extra_fields)) = (body.as_object_mut(), &self.extra_fields) {
+            for (k, v) in extra_fields {
+                if !body.contains_key(k) {
+                    body.insert(k.clone(), v.clone());
+                }
+            }
+        }
+    }
 }
 
 #[derive(Debug, Clone, Deserialize)]
 pub struct ModelConfig {
     pub name: String,
     pub max_tokens: Option<usize>,
+    pub extra_fields: Option<serde_json::Map<String, serde_json::Value>>,
     #[serde(deserialize_with = "deserialize_capabilities")]
     #[serde(default = "default_capabilities")]
     pub capabilities: ModelCapabilities,
-- 
cgit v1.2.3