Merge pull request #47 from VinciGit00/ollama-support

Ollama support

Author: VinciGit00

.streamlit/config.toml

@@ -0,0 +1,13 @@
+version: '3.8'
+services:
+  ollama:
+    image: ollama/ollama
+    container_name: ollama
+    ports:
+      - "11434:11434"
+    volumes:
+      - ollama_volume:/root/.ollama
+    restart: unless-stopped
+
+volumes:
+  ollama_volume:

.streamlit/credentials.toml

@@ -1,2 +1,4 @@
-OPENAI_APIKEY="your openai.com api key"
-GOOGLE_APIKEY="your google.com api key"
+OPENAI_APIKEY="your openai api key"
+GOOGLE_APIKEY="your google api key"
+AZURE_OPENAI_API_KEY="your azure api key"
+AZURE_OPENAI_ENDPOINT="https://<your-endpoint>.openai.azure.com/"

docker-compose.yml

@@ -7,11 +7,15 @@
 from scrapegraphai.models import Gemini
 from scrapegraphai.graphs import BaseGraph
 from scrapegraphai.nodes import FetchNode, ParseNode, RAGNode, GenerateAnswerNode
-
+from scrapegraphai.helpers import models_tokens
 load_dotenv()
-gemini_key = os.getenv("GOOGLE_APIKEY")
 
+# ************************************************
 # Define the configuration for the graph
+# ************************************************
+
+gemini_key = os.getenv("GOOGLE_APIKEY")
+
 graph_config = {
     "llm": {
         "api_key": gemini_key,
@@ -21,6 +25,10 @@
     },
 }
 
+# ************************************************
+# Define the graph nodes
+# ************************************************
+
 llm_model = Gemini(graph_config["llm"])
 
 # define the nodes for the graph
@@ -31,19 +39,23 @@
 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"],
-    model_config={"llm_model": llm_model},
+    node_config={"llm": llm_model},
 )
 generate_answer_node = GenerateAnswerNode(
     input="user_prompt & (relevant_chunks | parsed_doc | doc)",
     output=["answer"],
-    model_config={"llm_model": llm_model},
+    node_config={"llm": llm_model},
 )
 
-# create the graph by defining the nodes and their connections
+# ************************************************
+# Create the graph by defining the connections
+# ************************************************
+
 graph = BaseGraph(
     nodes={
         fetch_node,
@@ -59,7 +71,10 @@
     entry_point=fetch_node
 )
 
-# execute the graph
+# ************************************************
+# Execute the graph
+# ************************************************
+
 result = graph.execute({
     "user_prompt": "List me the projects with their description",
     "url": "https://perinim.github.io/projects/"

examples/.env.example

@@ -7,11 +7,15 @@
 from scrapegraphai.models import OpenAI
 from scrapegraphai.graphs import BaseGraph
 from scrapegraphai.nodes import FetchNode, ParseNode, RAGNode, GenerateAnswerNode
-
+from scrapegraphai.helpers import models_tokens
 load_dotenv()
-openai_key = os.getenv("OPENAI_APIKEY")
 
+# ************************************************
 # Define the configuration for the graph
+# ************************************************
+
+openai_key = os.getenv("OPENAI_APIKEY")
+
 graph_config = {
     "llm": {
         "api_key": openai_key,
@@ -21,6 +25,10 @@
     },
 }
 
+# ************************************************
+# Define the graph nodes
+# ************************************************
+
 llm_model = OpenAI(graph_config["llm"])
 
 # define the nodes for the graph
@@ -31,19 +39,23 @@
 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"],
-    model_config={"llm_model": llm_model},
+    node_config={"llm": llm_model},
 )
 generate_answer_node = GenerateAnswerNode(
     input="user_prompt & (relevant_chunks | parsed_doc | doc)",
     output=["answer"],
-    model_config={"llm_model": llm_model},
+    node_config={"llm": llm_model},
 )
 
-# create the graph by defining the nodes and their connections
+# ************************************************
+# Create the graph by defining the connections
+# ************************************************
+
 graph = BaseGraph(
     nodes={
         fetch_node,
@@ -59,7 +71,10 @@
     entry_point=fetch_node
 )
 
-# execute the graph
+# ************************************************
+# Execute the graph
+# ************************************************
+
 result = graph.execute({
     "user_prompt": "List me the projects with their description",
     "url": "https://perinim.github.io/projects/"

examples/custom_graph_gemini.py

@@ -6,26 +6,40 @@
 from dotenv import load_dotenv
 from scrapegraphai.graphs import SmartScraperGraph
 from scrapegraphai.utils import convert_to_csv, convert_to_json
-
 load_dotenv()
-openai_key = os.getenv("OPENAI_APIKEY")
 
+# ************************************************
+# Read the text file
+# ************************************************
+
+FILE_NAME = "inputs/plain_html_example.txt"
+curr_dir = os.path.dirname(os.path.realpath(__file__))
+file_path = os.path.join(curr_dir, FILE_NAME)
+
+# It could be also a http request using the request model
+with open(file_path, 'r', encoding="utf-8") as file:
+    text = file.read()
+
+# ************************************************
 # Define the configuration for the graph
+# ************************************************
+
+openai_key = os.getenv("OPENAI_APIKEY")
+
 graph_config = {
     "llm": {
         "api_key": openai_key,
         "model": "gpt-3.5-turbo",
     },
 }
 
+# ************************************************
+# Create the SmartScraperGraph instance and run it
+# ************************************************
 
-# It could be also a http request using the request model
-text = open('inputs/plain_html_example.txt', 'r', encoding="utf-8")
-
-# Create the SmartScraperGraph instance
 smart_scraper_graph = SmartScraperGraph(
     prompt="List me all the news with their description.",
-    file_source=str(text),
+    source=text,
     config=graph_config
 )
 

examples/custom_graph_openai.py

@@ -6,26 +6,39 @@
 from dotenv import load_dotenv
 from scrapegraphai.graphs import SmartScraperGraph
 from scrapegraphai.utils import convert_to_csv, convert_to_json
-
 load_dotenv()
-openai_key = os.getenv("OPENAI_APIKEY")
 
+# ************************************************
+# Read the XML file
+# ************************************************
+
+FILE_NAME = "inputs/books.xml"
+curr_dir = os.path.dirname(os.path.realpath(__file__))
+file_path = os.path.join(curr_dir, FILE_NAME)
+
+with open(file_path, 'r', encoding="utf-8") as file:
+    text = file.read()
+
+# ************************************************
 # Define the configuration for the graph
+# ************************************************
+
+openai_key = os.getenv("OPENAI_APIKEY")
+
 graph_config = {
     "llm": {
         "api_key": openai_key,
         "model": "gpt-3.5-turbo",
     },
 }
 
-# Read the XML file
-with open('inputs/books.xml', 'r', encoding="utf-8") as file:
-    text = file.read()
+# ************************************************
+# Create the SmartScraperGraph instance and run it
+# ************************************************
 
-# Create the SmartScraperGraph instance
 smart_scraper_graph = SmartScraperGraph(
     prompt="List me all the authors, title and genres of the books",
-    file_source=text,  # Pass the content of the file, not the file object
+    source=text,  # Pass the content of the file, not the file object
     config=graph_config
 )
 

examples/graph_evaluation_example.py

@@ -6,20 +6,30 @@
 from dotenv import load_dotenv
 from scrapegraphai.graphs import SearchGraph
 from scrapegraphai.utils import convert_to_csv, convert_to_json
-
 load_dotenv()
-openai_key = os.getenv("OPENAI_APIKEY")
 
+# ************************************************
 # Define the configuration for the graph
+# ************************************************
+
+openai_key = os.getenv("OPENAI_APIKEY")
+
 graph_config = {
     "llm": {
+        "model": "ollama/mistral",
+        "temperature": 0,
+        "format": "json", # Ollama needs the format to be specified explicitly
+    },
+    "embeddings": {
         "api_key": openai_key,
         "model": "gpt-3.5-turbo",
-        "temperature": 0,
     },
 }
 
-# Create the SmartScraperGraph instance
+# ************************************************
+# Create the SearchGraph instance and run it
+# ************************************************
+
 search_graph = SearchGraph(
     prompt="List me all the regions of Italy.",
     config=graph_config

examples/scrape_plain_text.py

@@ -5,23 +5,30 @@
 import os
 from dotenv import load_dotenv
 from scrapegraphai.graphs import SmartScraperGraph
-
 load_dotenv()
-openai_key = os.getenv("OPENAI_APIKEY")
 
+
+# ************************************************
 # Define the configuration for the graph
+# ************************************************
+
+openai_key = os.getenv("OPENAI_APIKEY")
+
 graph_config = {
     "llm": {
         "api_key": openai_key,
         "model": "gpt-3.5-turbo",
     },
 }
 
-# Create the SmartScraperGraph instance
+# ************************************************
+# Create the SmartScraperGraph instance and run it
+# ************************************************
+
 smart_scraper_graph = SmartScraperGraph(
     prompt="List me all the news with their description.",
     # also accepts a string with the already downloaded HTML code
-    file_source="https://www.wired.com",
+    source="https://www.wired.com",
     config=graph_config
 )