From de292b8fff0a4bab58e02afeaf845d9501a39562 Mon Sep 17 00:00:00 2001 From: Gabriel Luiz Freitas Almeida Date: Sun, 7 Jan 2024 20:06:20 -0300 Subject: [PATCH] Refactor OllamaLLM imports and remove unnecessary code --- src/backend/langflow/components/llms/OllamaLLM.py | 10 ++++------ 1 file changed, 4 insertions(+), 6 deletions(-) diff --git a/src/backend/langflow/components/llms/OllamaLLM.py b/src/backend/langflow/components/llms/OllamaLLM.py index 9bf00ad5d..abbc00387 100644 --- a/src/backend/langflow/components/llms/OllamaLLM.py +++ b/src/backend/langflow/components/llms/OllamaLLM.py @@ -1,8 +1,7 @@ -from typing import Optional, List +from typing import List, Optional -from langchain.llms import Ollama from langchain.llms.base import BaseLLM - +from langchain_community.llms.ollama import Ollama from langflow import CustomComponent @@ -150,10 +149,9 @@ class OllamaLLM(CustomComponent): "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} + # None Value remove + llm_params = {k: v for k, v in llm_params.items() if v is not None} try: llm = Ollama(**llm_params)