From c7d45b6d4bc33593e84213a16c85c5ce5c482bd9 Mon Sep 17 00:00:00 2001
From: Yilin Xu <xu.yilin@northeastern.edu>
Date: Tue, 30 Apr 2024 21:50:17 -0700
Subject: [PATCH] fix(dspy): added `copy` to `OllamaLocal` to propagate
 `model`, `model_type`, `base_url` and `timeout_s`.

---
 dsp/modules/ollama.py | 12 ++++++++++++
 1 file changed, 12 insertions(+)

diff --git a/dsp/modules/ollama.py b/dsp/modules/ollama.py
index 27304d271e..7781583bef 100644
--- a/dsp/modules/ollama.py
+++ b/dsp/modules/ollama.py
@@ -180,3 +180,15 @@ def __call__(
         completions = [self._get_choice_text(c) for c in choices]
 
         return completions
+    
+    def copy(self, **kwargs):
+        """Returns a copy of the language model with the same parameters."""
+        kwargs = {**self.kwargs, **kwargs}
+
+        return self.__class__(
+            model=self.model_name,
+            model_type=self.model_type,
+            base_url=self.base_url,
+            timeout_s=self.timeout_s,
+            **kwargs,
+        )