From 5f1e6bc6fa373ab61a6262f2e1c11c8c1db2b6f7 Mon Sep 17 00:00:00 2001
From: J S <49557684+svilupp@users.noreply.github.com>
Date: Sun, 3 Dec 2023 11:19:20 +0000
Subject: [PATCH] update docs + type

---
 CHANGELOG.md                             |  1 +
 docs/src/examples/working_with_ollama.md | 30 ++++++++++++++++++++++++
 src/llm_interface.jl                     |  2 +-
 3 files changed, 32 insertions(+), 1 deletion(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 1219d3438..1f67f8094 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -9,6 +9,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 ### Added
 
 ### Fixed
+- Changed type of global `PROMPT_SCHEMA::AbstractPromptSchema` for an easier switch to local models as a default option
 
 ## [0.2.0]
 
diff --git a/docs/src/examples/working_with_ollama.md b/docs/src/examples/working_with_ollama.md
index b4aeea77a..39be7036c 100644
--- a/docs/src/examples/working_with_ollama.md
+++ b/docs/src/examples/working_with_ollama.md
@@ -38,6 +38,36 @@ model = "openhermes2.5-mistral"
 "openhermes2.5-mistral"
 ````
 
+## Setting Ollama as a default LLM
+
+We need to change the global variables for PROMPT_SCHEMA and default models
+
+```julia
+using PromptingTools
+const PT = PromptingTools
+
+
+PT.PROMPT_SCHEMA = PT.OllamaManagedSchema()
+PT.MODEL_CHAT = "openhermes2.5-mistral"
+# You could do the same for PT.MODEL_EMBEDDING
+```
+
+We can also add a nicer alias for the above Mistral model
+
+```julia
+PT.MODEL_ALIASES["mistral"]= "openhermes2.5-mistral"
+# potentially also yi 34bn if you want a bigger more powerful model
+PT.MODEL_ALIASES["yi"]= "yi:34b-chat"
+```
+
+Now, we can use the `@ai_str` macro with Ollama models:
+```julia
+ai"Say hi to me!" # defaults to mistral because we set MODEL_CHAT above
+ai"Say hi to me in Chinese!"yi # defaults to yi 34Bn model
+```
+
+Note: Another quite popular model is `zephyr:7b-beta`
+
 ## Text Generation with aigenerate
 
 ### Simple message
diff --git a/src/llm_interface.jl b/src/llm_interface.jl
index 3921dcdee..0dc4952f9 100644
--- a/src/llm_interface.jl
+++ b/src/llm_interface.jl
@@ -79,7 +79,7 @@ struct OllamaManagedSchema <: AbstractOllamaManagedSchema end
 end
 
 ## Dispatch into default schema
-const PROMPT_SCHEMA = OpenAISchema()
+const PROMPT_SCHEMA::AbstractPromptSchema = OpenAISchema()
 
 aigenerate(prompt; kwargs...) = aigenerate(PROMPT_SCHEMA, prompt; kwargs...)
 function aiembed(doc_or_docs, args...; kwargs...)