Merge pull request #604 from dufeng1010/add_ollama_timeout

add ollama_timeout param
vanna-ai · Aug 16, 2024 · eaf3f5a · eaf3f5a
2 parents d39dae3 + 22aa8eb
commit eaf3f5a
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/src/vanna/ollama/ollama.py b/src/vanna/ollama/ollama.py
@@ -27,7 +27,9 @@ def __init__(self, config=None):
     if ":" not in self.model:
       self.model += ":latest"
 
-    self.ollama_client = ollama.Client(self.host, timeout=Timeout(240.0))
+    self.ollama_timeout = config.get("ollama_timeout", 240.0)
+
+    self.ollama_client = ollama.Client(self.host, timeout=Timeout(self.ollama_timeout))
     self.keep_alive = config.get('keep_alive', None)
     self.ollama_options = config.get('options', {})
     self.num_ctx = self.ollama_options.get('num_ctx', 2048)