added a function calling support? method

andthattoo · Sep 16, 2024 · 31b6da6 · 31b6da6
1 parent f36c04f
commit 31b6da6
Show file tree

Hide file tree

Showing 2 changed files with 20 additions and 2 deletions.
diff --git a/src/program/models.rs b/src/program/models.rs
@@ -66,9 +66,25 @@ pub enum Model {
     #[serde(rename = "o1-preview")]
     O1Preview,
 }
-// phi3.5:3.8b
 
-// phi3.5:3.8b-mini-instruct-fp16
+impl Model {
+    pub fn supports_tool_calling(&self) -> bool {
+        match self {
+            // OpenAI models that support tool calling
+            Model::GPT4Turbo | Model::GPT4o | Model::GPT4oMini => true,
+            // Ollama models that support tool calling
+            Model::Llama3_1_8B
+            | Model::Llama3_1_8Bq8
+            | Model::Llama3_1_8Bf16
+            | Model::Phi3Medium
+            | Model::Phi3Medium128k
+            | Model::Gemma2_9BFp16
+            | Model::Qwen2_7Bf16
+            | Model::Qwen2_7B => true,
+            _ => false,
+        }
+    }
+}
 
 impl From<Model> for String {
     fn from(model: Model) -> Self {

diff --git a/tests/run.rs b/tests/run.rs
@@ -67,6 +67,8 @@ mod simple_workflow_tests {
         SIMPLE_WORKFLOW_PATH,
         "How does reiki work?"
     );
+
+    workflow_test!(simple_o1, Model::O1Mini, SIMPLE_WORKFLOW_PATH);
 }
 
 mod insert_workflow_tests {