|
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