Merge pull request #407 from ScrapeGraphAI/404-split-unit-testing-from-src

Merge #404 (split unit testing from src)

Author: f-aguzzi

scrapegraphai/graphs/abstract_graph.py

@@ -327,6 +327,7 @@ def _create_embedder(self, embedder_config: dict) -> object:
                 raise KeyError("Model not supported") from exc
             return OllamaEmbeddings(**embedder_params)
         elif "hugging_face" in embedder_params["model"]:
+            embedder_params["model"] = "/".join(embedder_params["model"].split("/")[1:])
             try:
                 models_tokens["hugging_face"][embedder_params["model"]]
             except KeyError as exc:

scrapegraphai/helpers/models_tokens.py

@@ -1,7 +1,3 @@
-"""
-Models token
-"""
-
 models_tokens = {
     "openai": {
         "gpt-3.5-turbo-0125": 16385,
@@ -31,8 +27,8 @@
     },
     "gemini": {
         "gemini-pro": 128000,
-        "gemini-1.5-flash-latest":128000,
-        "gemini-1.5-pro-latest":128000,
+        "gemini-1.5-flash-latest": 128000,
+        "gemini-1.5-pro-latest": 128000,
         "models/embedding-001": 2048
     },
     "ollama": { "command-r": 12800, 
@@ -102,9 +98,8 @@
         "mistral.mistral-7b-instruct-v0:2": 32768,
         "mistral.mixtral-8x7b-instruct-v0:1": 32768,
         "mistral.mistral-large-2402-v1:0": 32768,
-		# Embedding models
-		"amazon.titan-embed-text-v1": 8000,
-		"amazon.titan-embed-text-v2:0": 8000,
+        "amazon.titan-embed-text-v1": 8000,
+        "amazon.titan-embed-text-v2:0": 8000,
         "cohere.embed-english-v3": 512,
         "cohere.embed-multilingual-v3": 512
     },
@@ -156,6 +151,6 @@
         "ernie-bot-2-base-zh": 4096,
         "ernie-bot-2-base-en": 4096,
         "ernie-bot-2-base-en-zh": 4096,
-        "ernie-bot-2-base-zh-en": 4096,
+        "ernie-bot-2-base-zh-en": 4096
     }
 }

scrapegraphai/utils/parse_state_keys.py

@@ -101,18 +101,3 @@ def evaluate_expression(expression):
             final_result.append(key)
 
     return final_result
-
-
-EXPRESSION = "user_input & (relevant_chunks | parsed_document | document)"
-state = {
-    "user_input": None,
-    "document": None,
-    "parsed_document": None,
-    "relevant_chunks": None,
-}
-
-try:
-    result = parse_expression(EXPRESSION, state)
-    print("Matched keys:", result)
-except ValueError as e:
-    print("Error:", e)

tests/utils/parse_state_keys_test.py

@@ -0,0 +1,21 @@
+"""
+Parse_state_key test module 
+"""
+import pytest
+from scrapegraphai.utils.parse_state_keys import parse_expression
+
+
+def test_parse_expression():
+    """Test parse_expression function."""
+    EXPRESSION = "user_input & (relevant_chunks | parsed_document | document)"
+    state = {
+        "user_input": None,
+        "document": None,
+        "parsed_document": None,
+        "relevant_chunks": None,
+    }
+    try:
+        result = parse_expression(EXPRESSION, state)
+        assert result != []
+    except ValueError as e:
+        assert "Error" in str(e)