Merge branch 'main' into ollama-preload

Author: duranbe

examples/rag_example.py

@@ -6,27 +6,37 @@
 load_dotenv()
 Settings.setup_logging()
 
-persist_directory = './defaultDb'
+persist_directory = "./defaultDb"
 model_embeddings = Settings.DEFAULT_EMBEDDINGS_MODEL
-model_name = 'llama3.1:8b'
+model_name = "llama3.1:8b"
 system_prompt_directory = Settings.DEFAULT_SYSTEM_PROMPT
 collection_name = str(uuid.uuid4())
 
-rag = Builder() \
-    .with_embeddings(Settings.HUGGINGFACE, model_name=model_embeddings) \
-    .with_vector_store(Settings.CHROMA, persist_directory=persist_directory, collection_name=collection_name) \
-    .with_llm(Settings.OLLAMA, model_name=model_name, system_prompt=system_prompt_directory) \
+rag = (
+    Builder()
+    .with_embeddings(Settings.HUGGINGFACE, model_name=model_embeddings)
+    .with_vector_store(
+        Settings.CHROMA,
+        persist_directory=persist_directory,
+        collection_name=collection_name,
+    )
+    .with_llm(
+        Settings.OLLAMA,
+        model_name=model_name,
+        system_prompt=system_prompt_directory,
+        options={"num_ctx": 8192},
+    )
     .build_rag(k=5)
+)
 
 rag.vector_store.ingest(
-                    data_path='./src/raglight',
-                    # ignore_folders=ignore_folders
-                )
+    data_path="./src/raglight",
+    # ignore_folders=ignore_folders
+)
 while True:
     query = input(">>> ")
-    if query == "quit" or query == "bye" : 
-        print('🤖 : See you soon 👋')
+    if query == "quit" or query == "bye":
+        print("🤖 : See you soon 👋")
         break
     response = rag.generate(query)
     print(response)
-

src/raglight/llm/ollama_model.py

@@ -7,6 +7,11 @@
 from json import dumps
 import logging
 
+# https://docs.ollama.com/context-length
+OLLAMA_DEFAULT_CONTEXT_SIZE = 4096
+OLLAMA_OPTION_CONTEXT_SIZE = "num_ctx"
+OLLAMA_WARNING_CONTEXT_SIZE = 0.80
+
 
 class OllamaModel(LLM):
     """
@@ -47,10 +52,15 @@ def __init__(
         self.headers = headers
         self.preload_model = preload_model
         self.options = options
+        self.max_context_size = (
+            self.options.get(OLLAMA_OPTION_CONTEXT_SIZE, OLLAMA_DEFAULT_CONTEXT_SIZE)
+            if self.options
+            else OLLAMA_DEFAULT_CONTEXT_SIZE
+        )
         super().__init__(model_name, system_prompt, system_prompt_file, self.api_base)
         logging.info(f"Using Ollama with {model_name} model 🤖")
         self.role: str = role
-
+        
     @override
     def load(self) -> Client:
         """
@@ -100,6 +110,14 @@ def generate(self, input: Dict[str, Any]) -> str:
             messages=messages,
             options=self.options,
         )
+
+        token_usage = response.eval_count + response.prompt_eval_count
+        if token_usage / self.max_context_size > OLLAMA_WARNING_CONTEXT_SIZE:
+            logging.warning(
+                f"Over {OLLAMA_WARNING_CONTEXT_SIZE * 100}% of context window reached, consider increasing it or reducing prompt size."
+                + f" Current usage : {token_usage}  out of {self.max_context_size} Tokens"
+            )
+
         return response.message.content
 
     @override

tests/tests_llm/test_ollama_model.py

@@ -21,7 +21,9 @@ def setUp(self):
             role="assistant",
             content="Machine learning (ML) is a subset of artificial intelligence",
         )
-        chat_response: ChatResponse = ChatResponse(message=message)
+        chat_response: ChatResponse = ChatResponse(
+            message=message, prompt_eval_count=200, eval_count=50
+        )
         mock_ollama_client.chat = MagicMock(return_value=chat_response)
         self.model.model = mock_ollama_client