Skip to content

Commit

Permalink
Refactor ChatOllamaEndpoint.py
Browse files Browse the repository at this point in the history
  • Loading branch information
ogabrielluiz committed Dec 29, 2023
1 parent 82a33ec commit d9fe8c8
Showing 1 changed file with 3 additions and 4 deletions.
7 changes: 3 additions & 4 deletions src/backend/langflow/components/llms/ChatOllamaEndpoint.py
Original file line number Diff line number Diff line change
@@ -1,13 +1,12 @@
from typing import Optional, List, Dict, Any
from langchain.chat_models.base import BaseChatModel
from typing import Any, Dict, List, Optional

# from langchain_community.chat_models import ChatOllama
from langchain.chat_models import ChatOllama
from langchain.chat_models.base import BaseChatModel

# from langchain.chat_models import ChatOllama
from langflow import CustomComponent


# whe When a callback component is added to Langflow, the comment must be uncommented.
# from langchain.callbacks.manager import CallbackManager

Expand Down Expand Up @@ -246,7 +245,7 @@ def build(
llm_params = {k: v for k, v in llm_params.items() if v is not None}

try:
output = ChatOllama(**llm_params)
output = ChatOllama(**llm_params) # type: ignore
except Exception as e:
raise ValueError("Could not initialize Ollama LLM.") from e

Expand Down

0 comments on commit d9fe8c8

Please sign in to comment.