Scrapegraph-ai/scrapegraphai/utils/test_node.py
2024-03-17 16:48:58 +01:00

54 lines
1.3 KiB
Python

from scrapegraphai.nodes import FetchNode, ParseNode, RAGNode, GenerateAnswerNode
import os
from dotenv import load_dotenv
from scrapegraphai.models import OpenAI
load_dotenv()
# Define the configuration for the language model
openai_key = os.getenv("OPENAI_APIKEY")
llm_config = {
"api_key": openai_key,
"model_name": "gpt-3.5-turbo",
"temperature": 0,
"streaming": True
}
llm_model = OpenAI(llm_config)
state = {
"user_prompt": "List me all the projects",
"url": "https://perinim.github.io/projects/",
}
fetch_node = FetchNode(
input="url | local_dir",
output=["doc"],
node_name="fetch_html"
)
updated_state = fetch_node.execute(state)
parse_node = ParseNode(
input="doc",
output=["parsed_doc"],
node_name="parse_document"
)
updated_state = parse_node.execute(updated_state)
rag_node = RAGNode(
input="user_prompt & (parsed_doc | doc)",
output=["relevant_chunks"],
model_config={"llm_model": llm_model},
node_name="rag_node"
)
updated_state = rag_node.execute(updated_state)
generate_answer_node = GenerateAnswerNode(
input="user_prompt & (relevant_chunks | parsed_doc | doc)",
output=["answer"],
model_config={"llm_model": llm_model},
node_name="generate_answer"
)
print(generate_answer_node.execute(updated_state))