Update OllamaLLM.py

src/backend/langflow/components/llms/OllamaLLM.py

@@ -14,143 +14,149 @@ class OllamaLLM(CustomComponent):
         return {
             "base_url": {
                 "display_name": "Base URL",
-                "info": "Endpoint of the Ollama API. Defaults to 'http://localhost:11434' if not specified."
+                "info": "Endpoint of the Ollama API. Defaults to 'http://localhost:11434' if not specified.",
             },
             "model": {
                 "display_name": "Model Name",
                 "value": "llama2",
-                "info": "Refer to https://ollama.ai/library for more models."
+                "info": "Refer to https://ollama.ai/library for more models.",
             },
             "temperature": {
                 "display_name": "Temperature",
                 "field_type": "float",
                 "value": 0.8,
-                "info": "Controls the creativity of model responses."
+                "info": "Controls the creativity of model responses.",
             },
-
             "mirostat": {
                 "display_name": "Mirostat",
                 "options": ["Disabled", "Mirostat", "Mirostat 2.0"],
                 "info": "Enable/disable Mirostat sampling for controlling perplexity.",
                 "value": "Disabled",
-                "advanced": True
+                "advanced": True,
             },
             "mirostat_eta": {
                 "display_name": "Mirostat Eta",
                 "field_type": "float",
                 "info": "Learning rate influencing the algorithm's response to feedback.",
-                "advanced": True
+                "advanced": True,
             },
-
-
             "mirostat_tau": {
                 "display_name": "Mirostat Tau",
                 "field_type": "float",
-                "value": 5.0,
                 "info": "Controls balance between coherence and diversity.",
-                "advanced": True
+                "advanced": True,
             },
             "num_ctx": {
                 "display_name": "Context Window Size",
                 "field_type": "int",
-                "value": 2048,
                 "info": "Size of the context window for generating the next token.",
-                "advanced": True
+                "advanced": True,
             },
             "num_gpu": {
                 "display_name": "Number of GPUs",
                 "field_type": "int",
                 "info": "Number of GPUs to use for computation.",
-                "advanced": True
+                "advanced": True,
             },
             "num_thread": {
                 "display_name": "Number of Threads",
                 "field_type": "int",
                 "info": "Number of threads to use during computation.",
-                "advanced": True
+                "advanced": True,
             },
             "repeat_last_n": {
                 "display_name": "Repeat Last N",
                 "field_type": "int",
-                "value": 64,
                 "info": "Sets how far back the model looks to prevent repetition.",
-                "advanced": True
+                "advanced": True,
             },
             "repeat_penalty": {
                 "display_name": "Repeat Penalty",
                 "field_type": "float",
-                "value": 1.1,
                 "info": "Penalty for repetitions in generated text.",
-                "advanced": True
+                "advanced": True,
             },
-
             "stop": {
                 "display_name": "Stop Tokens",
-
                 "info": "List of tokens to signal the model to stop generating text.",
-                "advanced": True
+                "advanced": True,
             },
             "tfs_z": {
                 "display_name": "TFS Z",
                 "field_type": "float",
-                "value": 1,
                 "info": "Tail free sampling to reduce impact of less probable tokens.",
-                "advanced": True
+                "advanced": True,
             },
             "top_k": {
                 "display_name": "Top K",
                 "field_type": "int",
-                "value": 40,
                 "info": "Limits token selection to top K for reducing nonsense generation.",
-                "advanced": True
+                "advanced": True,
             },
             "top_p": {
                 "display_name": "Top P",
                 "field_type": "int",
-                "value": 0.9,
                 "info": "Works with top-k to control diversity of generated text.",
-                "advanced": True
+                "advanced": True,
             },
         }
 
-    def build(self, base_url: Optional[str], model: str, mirostat: str, mirostat_eta: Optional[float],
+    def build(
-              mirostat_tau: Optional[float], num_ctx: Optional[int], num_gpu: Optional[int],
+        self,
-              num_thread: Optional[int], repeat_last_n: Optional[int], repeat_penalty: Optional[float],
+        base_url: Optional[str],
-              temperature: Optional[float], stop: Optional[List[str]], tfs_z: Optional[float],
+        model: str,
-              top_k: Optional[int], top_p: Optional[int]) -> BaseLLM:
+        temperature: Optional[float],
-
+        mirostat: Optional[str],
+        mirostat_eta: Optional[float] = None,
+        mirostat_tau: Optional[float] = None,
+        num_ctx: Optional[int] = None,
+        num_gpu: Optional[int] = None,
+        num_thread: Optional[int] = None,
+        repeat_last_n: Optional[int] = None,
+        repeat_penalty: Optional[float] = None,
+        stop: Optional[List[str]] = None,
+        tfs_z: Optional[float] = None,
+        top_k: Optional[int] = None,
+        top_p: Optional[int] = None,
+    ) -> BaseLLM:
         if not base_url:
             base_url = "http://localhost:11434"
 
-        mirostat_value = 0  # Default value for 'Disabled'
+        # Mapping mirostat settings to their corresponding values
+        mirostat_options = {"Mirostat": 1, "Mirostat 2.0": 2}
 
-        # Map the textual option to the corresponding integer
+        # Default to 0 for 'Disabled'
-        if mirostat == "Mirostat":
+        mirostat_value = mirostat_options.get(mirostat, 0)  # type: ignore
-            mirostat_value = 1
+
-        elif mirostat == "Mirostat 2.0":
+        # Set mirostat_eta and mirostat_tau to None if mirostat is disabled
-            mirostat_value = 2
+        if mirostat_value == 0:
+            mirostat_eta = None
+            mirostat_tau = None
+
+        llm_params = {
+            "base_url": base_url,
+            "model": model,
+            "mirostat": mirostat_value,
+            "mirostat_eta": mirostat_eta,
+            "mirostat_tau": mirostat_tau,
+            "num_ctx": num_ctx,
+            "num_gpu": num_gpu,
+            "num_thread": num_thread,
+            "repeat_last_n": repeat_last_n,
+            "repeat_penalty": repeat_penalty,
+            "temperature": temperature,
+            "stop": stop,
+            "tfs_z": tfs_z,
+            "top_k": top_k,
+            "top_p": top_p,
+        }
+        
+                # None Value remove
+        llm_params = {k: v for k, v in llm_params.items() if v is not None}
 
-        params = {k: v for k, v in {
-            'base_url': base_url,
-            'model': model,
-            'mirostat': mirostat_value,
-            'mirostat_eta': mirostat_eta,
-            'mirostat_tau': mirostat_tau,
-            'num_ctx': num_ctx,
-            'num_gpu': num_gpu,
-            'num_thread': num_thread,
-            'repeat_last_n': repeat_last_n,
-            'repeat_penalty': repeat_penalty,
-            'temperature': temperature,
-            'stop': stop,
-            'tfs_z': tfs_z,
-            'top_k': top_k,
-            'top_p': top_p,
-            'streaming' :"True"
-        }.items() if v is not None}
 
         try:
-            llm = Ollama(**params)
+            llm = Ollama(**llm_params)
         except Exception as e:
             raise ValueError("Could not connect to Ollama.") from e