fixed text input example

This commit is contained in:
Perinim 2024-03-12 12:22:47 +01:00
parent 55702b28a0
commit 308c2d28c8
2 changed files with 13 additions and 7 deletions

View File

@ -6,7 +6,7 @@ import os
from dotenv import load_dotenv
from scrapegraphai.models import OpenAI
from scrapegraphai.graphs import BaseGraph
from scrapegraphai.nodes import FetchTextNode, ParseTextNode, GenerateAnswerNode
from scrapegraphai.nodes import FetchTextNode, ParseNode, RAGNode, GenerateAnswerNode
load_dotenv()
@ -20,13 +20,17 @@ llm_config = {
}
model = OpenAI(llm_config)
with open("text_example.txt", "r", encoding="utf-8") as file:
curr_dir = os.path.dirname(__file__)
file_path = os.path.join(curr_dir, "text_example.txt")
with open(file_path, "r", encoding="utf-8") as file:
text = file.read()
# define the nodes for the graph
fetch_html_node = FetchTextNode("load_html")
parse_document_node = ParseTextNode("parse_document")
fetch_html_node = FetchTextNode("load_html_from_text")
parse_document_node = ParseNode(doc_type="text", chunks_size=4000, node_name="parse_document")
rag_node = RAGNode(model, "rag")
generate_answer_node = GenerateAnswerNode(model, "generate_answer")
# create the graph
@ -34,11 +38,13 @@ graph = BaseGraph(
nodes={
fetch_html_node,
parse_document_node,
rag_node,
generate_answer_node
},
edges={
(fetch_html_node, parse_document_node),
(parse_document_node, generate_answer_node)
(parse_document_node, rag_node),
(rag_node, generate_answer_node)
},
entry_point=fetch_html_node
)

View File

@ -2,7 +2,7 @@
Module for FetchTextNode
"""
from .base_node import BaseNode
from langchain_core.documents import Document
class FetchTextNode(BaseNode):
"""
@ -53,5 +53,5 @@ class FetchTextNode(BaseNode):
if 'text' not in state:
raise KeyError("The 'url' key is required to load the text.")
state["document"] = state["text"]
state["document"] = Document(page_content=state["text"])
return state