Merge pull request #40 from VinciGit00/First-Gemini-Support

.gitignore

@@ -27,3 +27,4 @@ venv/
 *.mp3
 *.sqlite
 examples/graph_examples/ScrapeGraphAI_generated_graph
+main.py

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

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(

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

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):

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={

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")
 

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

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)

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
+from .fetch_text_node import FetchTextNode
+from .generate_answer_node_vanilla import GenerateAnswerNodeVanilla

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

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

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={