diff --git a/.gitignore b/.gitignore index 45329238..26f73e8c 100644 --- a/.gitignore +++ b/.gitignore @@ -27,3 +27,4 @@ venv/ *.mp3 *.sqlite examples/graph_examples/ScrapeGraphAI_generated_graph +main.py diff --git a/examples/graph_examples/custom_graph_example.py b/examples/graph_examples/custom_graph_example.py index e5c4c4a3..c22cd369 100644 --- a/examples/graph_examples/custom_graph_example.py +++ b/examples/graph_examples/custom_graph_example.py @@ -4,47 +4,55 @@ Example of custom graph using existing nodes import os from dotenv import load_dotenv -from scrapegraphai.models import OpenAI +#from scrapegraphai.models import OpenAI +from scrapegraphai.models import Gemini from scrapegraphai.graphs import BaseGraph -from scrapegraphai.nodes import FetchHTMLNode, ParseNode, RAGNode, GenerateAnswerNode +from scrapegraphai.nodes import FetchHTMLNode, ParseNode, GenerateAnswerNodeVanilla + + load_dotenv() # Define the configuration for the language model -openai_key = os.getenv("OPENAI_APIKEY") +""" openai_key = os.getenv("OPENAI_APIKEY") llm_config = { "api_key": openai_key, "model_name": "gpt-3.5-turbo", "temperature": 0, "streaming": True } -model = OpenAI(llm_config) +model = OpenAI(llm_config) """ + +gemini_key = os.getenv("GOOGLE_API_KEY") +llm_config = { + "api_key": gemini_key, + "model_name": "gemini-pro", +} + +model = Gemini(llm_config) # define the nodes for the graph fetch_html_node = FetchHTMLNode("fetch_html") parse_document_node = ParseNode(doc_type="html", chunks_size=4000, node_name="parse_document") -rag_node = RAGNode(model, "rag") -generate_answer_node = GenerateAnswerNode(model, "generate_answer") +generate_answer_node = GenerateAnswerNodeVanilla(model, "generate_answer") # create the graph graph = BaseGraph( nodes={ fetch_html_node, parse_document_node, - rag_node, generate_answer_node }, edges={ (fetch_html_node, parse_document_node), - (parse_document_node, rag_node), - (rag_node, generate_answer_node) + (parse_document_node,generate_answer_node) }, entry_point=fetch_html_node ) # execute the graph inputs = {"user_input": "List me the projects with their description", - "url": "https://perinim.github.io/projects/"} + "url": "https://perinim.github.io/projects/"} result = graph.execute(inputs) # get the answer from the result diff --git a/examples/graph_examples/graph_from_text_example.py b/examples/graph_examples/graph_from_text_example.py index 0b6733e9..97444f09 100644 --- a/examples/graph_examples/graph_from_text_example.py +++ b/examples/graph_examples/graph_from_text_example.py @@ -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, ParseNode, RAGNode, GenerateAnswerNode +from scrapegraphai.nodes import FetchTextNode, ParseNode, RAGNode, GenerateAnswerNodeFromRag load_dotenv() @@ -32,7 +32,7 @@ fetch_text_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") +generate_answer_node = GenerateAnswerNodeFromRag(model, "generate_answer") # create the graph graph = BaseGraph( diff --git a/requirements.txt b/requirements.txt index 5a289fef..3f387f1a 100644 --- a/requirements.txt +++ b/requirements.txt @@ -2,6 +2,7 @@ langchain==0.1.6 langchain_community==0.0.19 langchain_core==0.1.22 langchain_openai==0.0.5 +langchain_google_genai==0.0.11 faiss-cpu==1.7.4 html2text==2020.1.16 beautifulsoup4==4.12.3 diff --git a/scrapegraphai/builders/graph_builder.py b/scrapegraphai/builders/graph_builder.py index 1da4c928..fe1942ed 100644 --- a/scrapegraphai/builders/graph_builder.py +++ b/scrapegraphai/builders/graph_builder.py @@ -4,7 +4,7 @@ Module for making the graph building import graphviz from langchain_core.prompts import ChatPromptTemplate from langchain.chains import create_extraction_chain -from ..models import OpenAI +from ..models import OpenAI, Gemini from ..helpers import nodes_metadata, graph_schema @@ -68,6 +68,13 @@ class GraphBuilder: llm_params = {**llm_defaults, **self.llm_config} if "api_key" not in llm_params: raise ValueError("LLM configuration must include an 'api_key'.") + + # select the model based on the model name + if "gpt-" in llm_params["model_name"]: + return OpenAI(llm_params) + elif "gemini" in llm_params["model_name"]: + return Gemini(llm_params) + return OpenAI(llm_params) def _generate_nodes_description(self): diff --git a/scrapegraphai/graphs/smart_scraper_graph.py b/scrapegraphai/graphs/smart_scraper_graph.py index fc25ad5a..ff03f7fc 100644 --- a/scrapegraphai/graphs/smart_scraper_graph.py +++ b/scrapegraphai/graphs/smart_scraper_graph.py @@ -7,7 +7,7 @@ from ..nodes import ( FetchHTMLNode, ParseNode, RAGNode, - GenerateAnswerNode + GenerateAnswerNodeFromRag ) class SmartScraperGraph: @@ -78,7 +78,7 @@ class SmartScraperGraph: fetch_html_node = FetchHTMLNode("fetch_html") parse_document_node = ParseNode(doc_type="html", chunks_size=4000, node_name="parse_document") rag_node = RAGNode(self.llm, "rag") - generate_answer_node = GenerateAnswerNode(self.llm, "generate_answer") + generate_answer_node = GenerateAnswerNodeFromRag(self.llm, "generate_answer") return BaseGraph( nodes={ diff --git a/scrapegraphai/graphs/speech_summary_graph.py b/scrapegraphai/graphs/speech_summary_graph.py index a23af88f..c121ae02 100644 --- a/scrapegraphai/graphs/speech_summary_graph.py +++ b/scrapegraphai/graphs/speech_summary_graph.py @@ -8,7 +8,7 @@ from ..nodes import ( FetchHTMLNode, ParseNode, RAGNode, - GenerateAnswerNode, + GenerateAnswerNodeFromRag, TextToSpeechNode, ) @@ -82,7 +82,7 @@ class SpeechSummaryGraph: fetch_html_node = FetchHTMLNode("fetch_html") parse_document_node = ParseNode(doc_type="html", chunks_size=4000, node_name="parse_document") rag_node = RAGNode(self.llm, "rag") - generate_answer_node = GenerateAnswerNode(self.llm, "generate_answer") + generate_answer_node = GenerateAnswerNodeFromRag(self.llm, "generate_answer") text_to_speech_node = TextToSpeechNode( self.text_to_speech_model, "text_to_speech") diff --git a/scrapegraphai/models/__init__.py b/scrapegraphai/models/__init__.py index f41bd57d..985a9153 100644 --- a/scrapegraphai/models/__init__.py +++ b/scrapegraphai/models/__init__.py @@ -5,3 +5,4 @@ from .openai import OpenAI from .openai_itt import OpenAIImageToText from .openai_tts import OpenAITextToSpeech +from .gemini import Gemini diff --git a/scrapegraphai/models/gemini.py b/scrapegraphai/models/gemini.py new file mode 100644 index 00000000..51be4fa2 --- /dev/null +++ b/scrapegraphai/models/gemini.py @@ -0,0 +1,19 @@ +from langchain_google_genai import ChatGoogleGenerativeAI + + +class Gemini(ChatGoogleGenerativeAI): + """Class for wrapping gemini module""" + + def __init__(self, llm_config: dict): + """ + A wrapper for the Gemini class that provides default configuration + and could be extended with additional methods if needed. + + Args: + llm_config (dict): Configuration parameters for the language model. + such as model="gemini-pro" and api_key + """ + # change the key model_name to model + llm_config["model"] = llm_config["model_name"] + # Initialize the superclass (ChatOpenAI) with provided config parameters + super().__init__(**llm_config) diff --git a/scrapegraphai/nodes/__init__.py b/scrapegraphai/nodes/__init__.py index d4af797c..03e027c5 100644 --- a/scrapegraphai/nodes/__init__.py +++ b/scrapegraphai/nodes/__init__.py @@ -4,9 +4,10 @@ __init__.py file for node folder from .fetch_html_node import FetchHTMLNode from .conditional_node import ConditionalNode from .get_probable_tags_node import GetProbableTagsNode -from .generate_answer_node import GenerateAnswerNode +from .generate_answer_node_from_rag import GenerateAnswerNodeFromRag from .parse_node import ParseNode from .rag_node import RAGNode from .text_to_speech_node import TextToSpeechNode from .image_to_text_node import ImageToTextNode -from .fetch_text_node import FetchTextNode \ No newline at end of file +from .fetch_text_node import FetchTextNode +from .generate_answer_node_vanilla import GenerateAnswerNodeVanilla \ No newline at end of file diff --git a/scrapegraphai/nodes/generate_answer_node.py b/scrapegraphai/nodes/generate_answer_node_from_rag.py similarity index 99% rename from scrapegraphai/nodes/generate_answer_node.py rename to scrapegraphai/nodes/generate_answer_node_from_rag.py index 3524e187..6aca5af0 100644 --- a/scrapegraphai/nodes/generate_answer_node.py +++ b/scrapegraphai/nodes/generate_answer_node_from_rag.py @@ -13,7 +13,7 @@ from langchain_core.runnables import RunnableParallel from .base_node import BaseNode -class GenerateAnswerNode(BaseNode): +class GenerateAnswerNodeFromRag(BaseNode): """ A node that generates an answer using a language model (LLM) based on the user's input and the content extracted from a webpage. It constructs a prompt from the user's input diff --git a/scrapegraphai/nodes/generate_answer_node_vanilla.py b/scrapegraphai/nodes/generate_answer_node_vanilla.py new file mode 100644 index 00000000..a7cd9aaa --- /dev/null +++ b/scrapegraphai/nodes/generate_answer_node_vanilla.py @@ -0,0 +1,103 @@ +""" +Module for generating the answer node +""" +# Imports from standard library +from tqdm import tqdm + +# Imports from Langchain +from langchain.prompts import PromptTemplate +from langchain_core.output_parsers import JsonOutputParser +from langchain_core.runnables import RunnableParallel + +# Imports from the library +from .base_node import BaseNode + + +class GenerateAnswerNodeVanilla(BaseNode): + """ + A node that generates an answer using a language model (LLM) based on the user's input + and the content extracted from a webpage. It constructs a prompt from the user's input + and the scraped content, feeds it to the LLM, and parses the LLM's response to produce + an answer. + + Attributes: + llm (ChatOpenAI): An instance of a language model client, configured for generating answers. + node_name (str): The unique identifier name for the node, defaulting + to "GenerateAnswerNode". + node_type (str): The type of the node, set to "node" indicating a + standard operational node. + + Args: + llm: An instance of the language model client (e.g., ChatOpenAI) used + for generating answers. + node_name (str, optional): The unique identifier name for the node. + Defaults to "GenerateAnswerNodeVanilla". + + Methods: + execute(state): Processes the input and document from the state to generate an answer, + updating the state with the generated answer under the 'answer' key. + """ + + def __init__(self, llm, node_name: str): + """ + Initializes the GenerateAnswerNode with a language model client and a node name. + Args: + llm (OpenAIImageToText): An instance of the OpenAIImageToText class. + node_name (str): name of the node + """ + super().__init__(node_name, "node") + self.llm = llm + + def execute(self, state: dict) -> dict: + """ + Generates an answer by constructing a prompt from the user's input and the scraped + content, querying the language model, and parsing its response. + + The method updates the state with the generated answer under the 'answer' key. + + Args: + state (dict): The current state of the graph, expected to contain 'user_input', + and optionally 'parsed_document' or 'relevant_chunks' within 'keys'. + + Returns: + dict: The updated state with the 'answer' key containing the generated answer. + + Raises: + KeyError: If 'user_input' or 'document' is not found in the state, indicating + that the necessary information for generating an answer is missing. + """ + + print("---GENERATING ANSWER---") + try: + user_input = state["user_input"] + document = state["document"][0] + except KeyError as e: + print(f"Error: {e} not found in state.") + raise + + context = document + + output_parser = JsonOutputParser() + format_instructions = output_parser.get_format_instructions() + + template_json = """You are a website scraper and you have just scraped the + following content from a website. + You are now asked to answer a question about the content you have scraped.\n {format_instructions} \n + This is the scraped text:\n + {context} \n + Question: {question} + """ + + # Merge the answers from the chunks + merge_prompt = PromptTemplate( + template=template_json, + input_variables=["context", "question"], + partial_variables={"format_instructions": format_instructions}, + ) + merge_chain = merge_prompt | self.llm | output_parser + answer = merge_chain.invoke( + {"context": context, "question": user_input}) + + # Update the state with the generated answer + state.update({"answer": answer}) + return state \ No newline at end of file diff --git a/tests/graphs/graph_from_text_test.py b/tests/graphs/graph_from_text_test.py index 5c4154a8..c222e223 100644 --- a/tests/graphs/graph_from_text_test.py +++ b/tests/graphs/graph_from_text_test.py @@ -6,7 +6,7 @@ import unittest from unittest.mock import patch from scrapegraphai.models import OpenAI from scrapegraphai.graphs import BaseGraph -from scrapegraphai.nodes import FetchTextNode, ParseNode, RAGNode, GenerateAnswerNode +from scrapegraphai.nodes import FetchTextNode, ParseNode, RAGNode, GenerateAnswerNodeFromRag class TestCustomGraph(unittest.TestCase): @@ -59,7 +59,7 @@ class TestCustomGraph(unittest.TestCase): parse_document_node = ParseNode( doc_type="text", chunks_size=20, node_name="parse_document") rag_node = RAGNode(model, "rag") - generate_answer_node = GenerateAnswerNode(model, "generate_answer") + generate_answer_node = GenerateAnswerNodeFromRag(model, "generate_answer") graph = BaseGraph( nodes={