feat: add integration for chuinking for openai

VinciGit00 · VinciGit00 · commit f5a1dbb2a557 · 2024-08-24T16:53:44.000+02:00
diff --git a/examples/anthropic/custom_graph_haiku.py b/examples/anthropic/custom_graph_haiku.py
@@ -27,7 +27,7 @@
 # Define the graph nodes
 # ************************************************
 
-llm_model = OpenAI(graph_config["llm"])
+llm_model = ChatOpenAI(graph_config["llm"])
 embedder = OpenAIEmbeddings(api_key=llm_model.openai_api_key)
 
 # define the nodes for the graph
@@ -55,15 +55,7 @@
     node_config={
         "chunk_size": 4096,
         "verbose": True,
-    }
-)
-rag_node = RAGNode(
-    input="user_prompt & (parsed_doc | doc)",
-    output=["relevant_chunks"],
-    node_config={
         "llm_model": llm_model,
-        "embedder_model": embedder,
-        "verbose": True,
     }
 )
 generate_answer_node = GenerateAnswerNode(
diff --git a/examples/ernie/custom_graph_ernie.py b/examples/ernie/custom_graph_ernie.py
@@ -4,11 +4,10 @@
 
 import os
 from dotenv import load_dotenv
-
 from langchain_openai import OpenAIEmbeddings
 from langchain_openai import ChatOpenAI
 from scrapegraphai.graphs import BaseGraph
-from scrapegraphai.nodes import FetchNode, ParseNode, RAGNode, GenerateAnswerNode, RobotsNode
+from scrapegraphai.nodes import FetchNode, ParseNode, GenerateAnswerNode, RobotsNode
 
 # ************************************************
 # Define the configuration for the graph
@@ -31,8 +30,7 @@
 # Define the graph nodes
 # ************************************************
 
-llm_model = OpenAI(graph_config["llm"])
-embedder = OpenAIEmbeddings(api_key=llm_model.openai_api_key)
+llm_model = ChatOpenAI(graph_config["llm"])
 
 # define the nodes for the graph
 robot_node = RobotsNode(
@@ -59,17 +57,10 @@
     node_config={
         "chunk_size": 4096,
         "verbose": True,
-    }
-)
-rag_node = RAGNode(
-    input="user_prompt & (parsed_doc | doc)",
-    output=["relevant_chunks"],
-    node_config={
         "llm_model": llm_model,
-        "embedder_model": embedder,
-        "verbose": True,
     }
 )
+
 generate_answer_node = GenerateAnswerNode(
     input="user_prompt & (relevant_chunks | parsed_doc | doc)",
     output=["answer"],
@@ -88,14 +79,12 @@
         robot_node,
         fetch_node,
         parse_node,
-        rag_node,
         generate_answer_node,
     ],
     edges=[
         (robot_node, fetch_node),
         (fetch_node, parse_node),
-        (parse_node, rag_node),
-        (rag_node, generate_answer_node)
+        (parse_node,  generate_answer_node)
     ],
     entry_point=robot_node
 )
diff --git a/examples/fireworks/custom_graph_fireworks.py b/examples/fireworks/custom_graph_fireworks.py
@@ -55,6 +55,7 @@
     node_config={
         "chunk_size": 4096,
         "verbose": True,
+        "llm_model": llm_model,
     }
 )
 
diff --git a/examples/google_genai/custom_graph_gemini.py b/examples/google_genai/custom_graph_gemini.py
@@ -38,13 +38,12 @@
 parse_node = ParseNode(
     input="doc",
     output=["parsed_doc"],
-    node_config={"chunk_size": 4096}
-)
-rag_node = RAGNode(
-    input="user_prompt & (parsed_doc | doc)",
-    output=["relevant_chunks"],
-    node_config={"llm": llm_model},
+    node_config={
+        "chunk_size": 4096,
+        "llm_model": llm_model,
+    }
 )
+
 generate_answer_node = GenerateAnswerNode(
     input="user_prompt & (relevant_chunks | parsed_doc | doc)",
     output=["answer"],
@@ -59,13 +58,11 @@
     nodes={
         fetch_node,
         parse_node,
-        rag_node,
         generate_answer_node,
     },
     edges={
         (fetch_node, parse_node),
-        (parse_node, rag_node),
-        (rag_node, generate_answer_node)
+        (parse_node, generate_answer_node)
     },
     entry_point=fetch_node
 )
diff --git a/examples/google_vertexai/custom_graph_gemini.py b/examples/google_vertexai/custom_graph_gemini.py
@@ -38,13 +38,10 @@
 parse_node = ParseNode(
     input="doc",
     output=["parsed_doc"],
-    node_config={"chunk_size": 4096}
-)
-rag_node = RAGNode(
-    input="user_prompt & (parsed_doc | doc)",
-    output=["relevant_chunks"],
-    node_config={"llm": llm_model},
+    node_config={"chunk_size": 4096,
+                         "llm_model": llm_model}
 )
+
 generate_answer_node = GenerateAnswerNode(
     input="user_prompt & (relevant_chunks | parsed_doc | doc)",
     output=["answer"],
diff --git a/examples/groq/custom_graph_groq.py b/examples/groq/custom_graph_groq.py
@@ -55,6 +55,7 @@
     node_config={
         "chunk_size": 4096,
         "verbose": True,
+        "llm_model": llm_model,
     }
 )
 
diff --git a/examples/huggingfacehub/custom_graph_huggingfacehub.py b/examples/huggingfacehub/custom_graph_huggingfacehub.py
@@ -67,17 +67,10 @@
     node_config={
         "chunk_size": 4096,
         "verbose": True,
-    }
-)
-rag_node = RAGNode(
-    input="user_prompt & (parsed_doc | doc)",
-    output=["relevant_chunks"],
-    node_config={
         "llm_model": llm_model,
-        "embedder_model": embedder,
-        "verbose": True,
     }
 )
+
 generate_answer_node = GenerateAnswerNode(
     input="user_prompt & (relevant_chunks | parsed_doc | doc)",
     output=["answer"],
diff --git a/examples/local_models/custom_graph_ollama.py b/examples/local_models/custom_graph_ollama.py
@@ -56,6 +56,7 @@
     node_config={
         "chunk_size": 4096,
         "verbose": True,
+        "llm_model": llm_model,
     }
 )
 
diff --git a/examples/mistral/custom_graph_mistral.py b/examples/mistral/custom_graph_mistral.py
@@ -54,6 +54,7 @@
     node_config={
         "chunk_size": 4096,
         "verbose": True,
+        "llm_model": llm_model,
     }
 )
 rag_node = RAGNode(
diff --git a/examples/mixed_models/custom_graph_groq_openai.py b/examples/mixed_models/custom_graph_groq_openai.py
@@ -63,6 +63,7 @@
     node_config={
         "chunk_size": 4096,
         "verbose": True,
+        "llm_model": llm_model,
     }
 )
 rag_node = RAGNode(
diff --git a/examples/nemotron/custom_graph_nemotron.py b/examples/nemotron/custom_graph_nemotron.py
@@ -4,7 +4,6 @@
 
 import os
 from dotenv import load_dotenv
-
 from langchain_openai import OpenAIEmbeddings
 from langchain_openai import ChatOpenAI
 from scrapegraphai.graphs import BaseGraph
@@ -26,8 +25,7 @@
 # Define the graph nodes
 # ************************************************
 
-llm_model = OpenAI(graph_config["llm"])
-embedder = OpenAIEmbeddings(api_key=llm_model.openai_api_key)
+llm_model = ChatOpenAI(graph_config["llm"])
 
 # define the nodes for the graph
 robot_node = RobotsNode(
@@ -54,15 +52,7 @@
     node_config={
         "chunk_size": 4096,
         "verbose": True,
-    }
-)
-rag_node = RAGNode(
-    input="user_prompt & (parsed_doc | doc)",
-    output=["relevant_chunks"],
-    node_config={
         "llm_model": llm_model,
-        "embedder_model": embedder,
-        "verbose": True,
     }
 )
 generate_answer_node = GenerateAnswerNode(
@@ -83,14 +73,12 @@
         robot_node,
         fetch_node,
         parse_node,
-        rag_node,
         generate_answer_node,
     ],
     edges=[
         (robot_node, fetch_node),
         (fetch_node, parse_node),
-        (parse_node, rag_node),
-        (rag_node, generate_answer_node)
+        (parse_node, generate_answer_node)
     ],
     entry_point=robot_node
 )
diff --git a/examples/oneapi/custom_graph_oneapi.py b/examples/oneapi/custom_graph_oneapi.py
@@ -22,8 +22,7 @@
 # Define the graph nodes
 # ************************************************
 
-llm_model = OpenAI(graph_config["llm"])
-embedder = OpenAIEmbeddings(api_key=llm_model.openai_api_key)
+llm_model = ChatOpenAI(graph_config["llm"])
 
 # define the nodes for the graph
 robot_node = RobotsNode(
@@ -50,15 +49,7 @@
     node_config={
         "chunk_size": 4096,
         "verbose": True,
-    }
-)
-rag_node = RAGNode(
-    input="user_prompt & (parsed_doc | doc)",
-    output=["relevant_chunks"],
-    node_config={
         "llm_model": llm_model,
-        "embedder_model": embedder,
-        "verbose": True,
     }
 )
 generate_answer_node = GenerateAnswerNode(
@@ -79,14 +70,12 @@
         robot_node,
         fetch_node,
         parse_node,
-        rag_node,
         generate_answer_node,
     ],
     edges=[
         (robot_node, fetch_node),
         (fetch_node, parse_node),
-        (parse_node, rag_node),
-        (rag_node, generate_answer_node)
+        (parse_node,generate_answer_node)
     ],
     entry_point=robot_node
 )
diff --git a/examples/openai/custom_graph_openai.py b/examples/openai/custom_graph_openai.py
@@ -27,8 +27,7 @@
 # Define the graph nodes
 # ************************************************
 
-llm_model = OpenAI(graph_config["llm"])
-embedder = OpenAIEmbeddings(api_key=llm_model.openai_api_key)
+llm_model = ChatOpenAI(graph_config["llm"])
 
 # define the nodes for the graph
 robot_node = RobotsNode(
@@ -55,17 +54,10 @@
     node_config={
         "chunk_size": 4096,
         "verbose": True,
-    }
-)
-rag_node = RAGNode(
-    input="user_prompt & (parsed_doc | doc)",
-    output=["relevant_chunks"],
-    node_config={
         "llm_model": llm_model,
-        "embedder_model": embedder,
-        "verbose": True,
     }
 )
+
 generate_answer_node = GenerateAnswerNode(
     input="user_prompt & (relevant_chunks | parsed_doc | doc)",
     output=["answer"],
@@ -84,14 +76,12 @@
         robot_node,
         fetch_node,
         parse_node,
-        rag_node,
         generate_answer_node,
     ],
     edges=[
         (robot_node, fetch_node),
         (fetch_node, parse_node),
-        (parse_node, rag_node),
-        (rag_node, generate_answer_node)
+        (parse_node,generate_answer_node)
     ],
     entry_point=robot_node
 )
diff --git a/pyproject.toml b/pyproject.toml
@@ -13,6 +13,8 @@ dependencies = [
     "langchain-google-genai>=1.0.7",
     "langchain-openai>=0.1.22",
     "langchain-mistralai>=0.1.12",
+    "mistral-common==1.3.4",
+    "langchain-ollama==0.1.1",
     "langchain_community>=0.2.9",
     "langchain-aws>=0.1.3",
     "html2text>=2024.2.26",
diff --git a/scrapegraphai/graphs/deep_scraper_graph.py b/scrapegraphai/graphs/deep_scraper_graph.py
@@ -75,6 +75,7 @@ def _create_repeated_graph(self) -> BaseGraph:
             input="doc",
             output=["parsed_doc"],
             node_config={
+                "llm_model": self.llm_model,
                 "chunk_size": self.model_token
             }
         )
diff --git a/scrapegraphai/graphs/markdown_scraper_graph.py b/scrapegraphai/graphs/markdown_scraper_graph.py
@@ -59,6 +59,7 @@ def _create_graph(self) -> BaseGraph:
             input="doc",
             output=["parsed_doc"],
             node_config={
+                "llm_model": self.llm_model,
                 "parse_html": False,
                 "chunk_size": self.model_token
             }
diff --git a/scrapegraphai/graphs/omni_scraper_graph.py b/scrapegraphai/graphs/omni_scraper_graph.py
@@ -74,6 +74,7 @@ def _create_graph(self) -> BaseGraph:
             input="doc",
             output=["parsed_doc"],
             node_config={
+                "llm_model": self.llm_model,
                 "chunk_size": self.model_token
             }
         )
diff --git a/scrapegraphai/graphs/script_creator_graph.py b/scrapegraphai/graphs/script_creator_graph.py
@@ -72,9 +72,11 @@ def _create_graph(self) -> BaseGraph:
         parse_node = ParseNode(
             input="doc",
             output=["parsed_doc"],
-            node_config={"chunk_size": self.model_token,
-                         "parse_html": False
-                         }
+            node_config = {
+                "llm_model": self.llm_model,
+                "chunk_size": self.model_token,
+                "parse_html": False
+            }
         )
         generate_scraper_node = GenerateScraperNode(
             input="user_prompt & (doc)",
diff --git a/scrapegraphai/graphs/search_link_graph.py b/scrapegraphai/graphs/search_link_graph.py
@@ -64,6 +64,7 @@ def _create_graph(self) -> BaseGraph:
             input="doc",
             output=["parsed_doc"],
             node_config={
+                "llm_model": self.llm_model,
                 "chunk_size": self.model_token
             }
         )
diff --git a/scrapegraphai/graphs/smart_scraper_graph.py b/scrapegraphai/graphs/smart_scraper_graph.py
@@ -74,6 +74,7 @@ def _create_graph(self) -> BaseGraph:
             input="doc",
             output=["parsed_doc"],
             node_config={
+                "llm_model": self.llm_model,
                 "chunk_size": self.model_token
             }
         )
diff --git a/scrapegraphai/graphs/speech_graph.py b/scrapegraphai/graphs/speech_graph.py
diff --git a/scrapegraphai/nodes/parse_node.py b/scrapegraphai/nodes/parse_node.py
diff --git a/scrapegraphai/utils/__init__.py b/scrapegraphai/utils/__init__.py
diff --git a/scrapegraphai/utils/tokenizers/tokenizer_mistral.py b/scrapegraphai/utils/tokenizers/tokenizer_mistral.py
diff --git a/scrapegraphai/utils/tokenizers/tokenizer_openai.py b/scrapegraphai/utils/tokenizers/tokenizer_openai.py

Original file line number	Diff line number	Diff line change
`@@ -55,6 +55,7 @@`
`55`	`55`	`node_config={`
`56`	`56`	`"chunk_size": 4096,`
`57`	`57`	`"verbose": True,`
	`58`	`+ "llm_model": llm_model,`
`58`	`59`	`}`
`59`	`60`	`)`
`60`	`61`
Original file line number	Diff line number	Diff line change
`@@ -56,6 +56,7 @@`
`56`	`56`	`node_config={`
`57`	`57`	`"chunk_size": 4096,`
`58`	`58`	`"verbose": True,`
	`59`	`+ "llm_model": llm_model,`
`59`	`60`	`}`
`60`	`61`	`)`
`61`	`62`
Original file line number	Diff line number	Diff line change
`@@ -54,6 +54,7 @@`
`54`	`54`	`node_config={`
`55`	`55`	`"chunk_size": 4096,`
`56`	`56`	`"verbose": True,`
	`57`	`+ "llm_model": llm_model,`
`57`	`58`	`}`
`58`	`59`	`)`
`59`	`60`	`rag_node = RAGNode(`
Original file line number	Diff line number	Diff line change
`@@ -63,6 +63,7 @@`
`63`	`63`	`node_config={`
`64`	`64`	`"chunk_size": 4096,`
`65`	`65`	`"verbose": True,`
	`66`	`+ "llm_model": llm_model,`
`66`	`67`	`}`
`67`	`68`	`)`
`68`	`69`	`rag_node = RAGNode(`