|
6 | 6 | from langchain.prompts import PromptTemplate |
7 | 7 | from langchain_core.output_parsers import JsonOutputParser |
8 | 8 | from langchain_core.runnables import RunnableParallel |
| 9 | +from langchain_openai import ChatOpenAI |
| 10 | +from langchain_community.chat_models import ChatOllama |
9 | 11 | from tqdm import tqdm |
10 | 12 | from ..utils.logging import get_logger |
11 | | -from ..models import Ollama, OpenAI |
12 | 13 | from .base_node import BaseNode |
13 | 14 | from ..helpers import template_chunks, template_no_chunks, template_merge, template_chunks_md, template_no_chunks_md, template_merge_md |
14 | 15 |
|
@@ -41,7 +42,7 @@ def __init__( |
41 | 42 |
|
42 | 43 | self.llm_model = node_config["llm_model"] |
43 | 44 |
|
44 | | - if isinstance(node_config["llm_model"], Ollama): |
| 45 | + if isinstance(node_config["llm_model"], ChatOllama): |
45 | 46 | self.llm_model.format="json" |
46 | 47 |
|
47 | 48 | self.verbose = ( |
@@ -93,7 +94,7 @@ def execute(self, state: dict) -> dict: |
93 | 94 |
|
94 | 95 | format_instructions = output_parser.get_format_instructions() |
95 | 96 |
|
96 | | - if isinstance(self.llm_model, OpenAI) and not self.script_creator or self.force and not self.script_creator or self.is_md_scraper: |
| 97 | + if isinstance(self.llm_model, ChatOpenAI) and not self.script_creator or self.force and not self.script_creator or self.is_md_scraper: |
97 | 98 | template_no_chunks_prompt = template_no_chunks_md |
98 | 99 | template_chunks_prompt = template_chunks_md |
99 | 100 | template_merge_prompt = template_merge_md |
|
0 commit comments