Align knowledge-graph and knowledge-store linting (#540)

cbornet · web-flow · commit 9fe63e1c0963 · 2024-06-28T17:19:11.000+02:00
diff --git a/libs/knowledge-graph/notebooks/notebook.ipynb b/libs/knowledge-graph/notebooks/notebook.ipynb
@@ -55,7 +55,9 @@
     "\n",
     "os.environ[\"OPENAI_API_KEY\"] = getpass.getpass(\"Enter OpenAI API Key: \")\n",
     "os.environ[\"ASTRA_DB_DATABASE_ID\"] = input(\"Enter Astra DB Database ID: \")\n",
-    "os.environ[\"ASTRA_DB_APPLICATION_TOKEN\"] = getpass.getpass(\"Enter Astra DB Application Token: \")\n",
+    "os.environ[\"ASTRA_DB_APPLICATION_TOKEN\"] = getpass.getpass(\n",
+    "    \"Enter Astra DB Application Token: \"\n",
+    ")\n",
     "\n",
     "keyspace = input(\"Enter Astra DB Keyspace (Empty for default): \")\n",
     "if keyspace:\n",
@@ -75,8 +77,8 @@
      "text": [
       "Requirement already satisfied: python-dotenv in /Users/benjamin.chambers/Library/Caches/pypoetry/virtualenvs/knowledge-graph-bxUBmW8M-py3.11/lib/python3.11/site-packages (1.0.1)\n",
       "\n",
-      "\u001B[1m[\u001B[0m\u001B[34;49mnotice\u001B[0m\u001B[1;39;49m]\u001B[0m\u001B[39;49m A new release of pip is available: \u001B[0m\u001B[31;49m23.3.1\u001B[0m\u001B[39;49m -> \u001B[0m\u001B[32;49m24.0\u001B[0m\n",
-      "\u001B[1m[\u001B[0m\u001B[34;49mnotice\u001B[0m\u001B[1;39;49m]\u001B[0m\u001B[39;49m To update, run: \u001B[0m\u001B[32;49mpip install --upgrade pip\u001B[0m\n",
+      "\u001b[1m[\u001b[0m\u001b[34;49mnotice\u001b[0m\u001b[1;39;49m]\u001b[0m\u001b[39;49m A new release of pip is available: \u001b[0m\u001b[31;49m23.3.1\u001b[0m\u001b[39;49m -> \u001b[0m\u001b[32;49m24.0\u001b[0m\n",
+      "\u001b[1m[\u001b[0m\u001b[34;49mnotice\u001b[0m\u001b[1;39;49m]\u001b[0m\u001b[39;49m To update, run: \u001b[0m\u001b[32;49mpip install --upgrade pip\u001b[0m\n",
       "Note: you may need to restart the kernel to use updated packages.\n"
      ]
     },
@@ -96,6 +98,7 @@
     "# See `env.template` for an example of what you should have there.\n",
     "%pip install python-dotenv\n",
     "import dotenv\n",
+    "\n",
     "dotenv.load_dotenv()"
    ]
   },
@@ -114,6 +117,7 @@
    "source": [
     "# Initialize cassandra connection from environment variables).\n",
     "import cassio\n",
+    "\n",
     "cassio.init(auto=True)"
    ]
   },
@@ -132,6 +136,7 @@
    "source": [
     "# Create graph store.\n",
     "from ragstack_knowledge_graph.cassandra_graph_store import CassandraGraphStore\n",
+    "\n",
     "graph_store = CassandraGraphStore()"
    ]
   },
@@ -342,6 +347,7 @@
    "source": [
     "# Render the extracted graph to GraphViz.\n",
     "from ragstack_knowledge_graph.render import render_graph_documents\n",
+    "\n",
     "render_graph_documents(graph_documents)"
    ]
   },
@@ -450,7 +456,8 @@
    "source": [
     "# Example showing extracted entities (nodes)\n",
     "from ragstack_knowledge_graph import extract_entities\n",
-    "extract_entities(llm).invoke({ \"question\": \"Who is Marie Curie?\"})"
+    "\n",
+    "extract_entities(llm).invoke({\"question\": \"Who is Marie Curie?\"})"
    ]
   },
   {
@@ -474,15 +481,19 @@
    "outputs": [],
    "source": [
     "from operator import itemgetter\n",
-    "from langchain_core.runnables import RunnableLambda, RunnablePassthrough\n",
+    "\n",
     "from langchain_core.prompts import ChatPromptTemplate\n",
-    "from ragstack_knowledge_graph import extract_entities\n",
+    "from langchain_core.runnables import RunnableLambda, RunnablePassthrough\n",
     "from langchain_openai import ChatOpenAI\n",
-    "llm = ChatOpenAI(model_name = \"gpt-4\")\n",
+    "from ragstack_knowledge_graph import extract_entities\n",
+    "\n",
+    "llm = ChatOpenAI(model_name=\"gpt-4\")\n",
+    "\n",
     "\n",
     "def _combine_relations(relations):\n",
     "    return \"\\n\".join(map(repr, relations))\n",
     "\n",
+    "\n",
     "ANSWER_PROMPT = (\n",
     "    \"The original question is given below.\"\n",
     "    \"This question has been used to retrieve information from a knowledge graph.\"\n",
@@ -494,10 +505,14 @@
     ")\n",
     "\n",
     "chain = (\n",
-    "    { \"question\": RunnablePassthrough() }\n",
-    "    | RunnablePassthrough.assign(entities = extract_entities(llm))\n",
-    "    | RunnablePassthrough.assign(triples = itemgetter(\"entities\") | graph_store.as_runnable())\n",
-    "    | RunnablePassthrough.assign(context = itemgetter(\"triples\") | RunnableLambda(_combine_relations))\n",
+    "    {\"question\": RunnablePassthrough()}\n",
+    "    | RunnablePassthrough.assign(entities=extract_entities(llm))\n",
+    "    | RunnablePassthrough.assign(\n",
+    "        triples=itemgetter(\"entities\") | graph_store.as_runnable()\n",
+    "    )\n",
+    "    | RunnablePassthrough.assign(\n",
+    "        context=itemgetter(\"triples\") | RunnableLambda(_combine_relations)\n",
+    "    )\n",
     "    | ChatPromptTemplate.from_messages([ANSWER_PROMPT])\n",
     "    | llm\n",
     ")"
diff --git a/libs/knowledge-graph/pyproject.toml b/libs/knowledge-graph/pyproject.toml
@@ -40,26 +40,6 @@ setuptools = "^70.0.0"
 requires = ["poetry-core"]
 build-backend = "poetry.core.masonry.api"
 
-[tool.ruff.lint]
-# Enable Pyflakes `E` and `F` codes by default.
-select = [
-    # Pyflakes
-    "F",
-    # Pycodestyle
-    "E",
-    "W",
-    # isort
-    "I001",
-]
-ignore = []
-
-# Allow unused variables when underscore-prefixed.
-dummy-variable-rgx = "^(_+|(_+[a-zA-Z0-9_]*[a-zA-Z0-9]+?))$"
-
-[tool.ruff.lint.mccabe]
-# Unlike Flake8, default to a complexity level of 10.
-max-complexity = 10
-
 [tool.mypy]
 strict = true
 warn_unreachable = true
diff --git a/libs/knowledge-graph/tests/conftest.py b/libs/knowledge-graph/tests/conftest.py
@@ -7,11 +7,10 @@
 from langchain.graphs.graph_document import GraphDocument, Node, Relationship
 from langchain_core.documents import Document
 from langchain_core.language_models import BaseChatModel
+from ragstack_knowledge_graph.cassandra_graph_store import CassandraGraphStore
 from testcontainers.core.container import DockerContainer
 from testcontainers.core.waiting_utils import wait_for_logs
 
-from ragstack_knowledge_graph.cassandra_graph_store import CassandraGraphStore
-
 load_dotenv()
 
 
diff --git a/libs/knowledge-graph/tests/test_extraction.py b/libs/knowledge-graph/tests/test_extraction.py
@@ -4,7 +4,6 @@
 from langchain_community.graphs.graph_document import Node, Relationship
 from langchain_core.documents import Document
 from langchain_core.language_models import BaseChatModel
-
 from ragstack_knowledge_graph.extraction import (
     KnowledgeSchema,
     KnowledgeSchemaExtractor,
diff --git a/libs/knowledge-graph/tests/test_knowledge_graph.py b/libs/knowledge-graph/tests/test_knowledge_graph.py
@@ -2,7 +2,6 @@
 
 import pytest
 from cassandra.cluster import Session
-
 from ragstack_knowledge_graph.knowledge_graph import CassandraKnowledgeGraph
 from ragstack_knowledge_graph.traverse import Node, Relation
 
diff --git a/libs/knowledge-graph/tests/test_schema_inference.py b/libs/knowledge-graph/tests/test_schema_inference.py
@@ -3,7 +3,6 @@
 import pytest
 from langchain_core.documents import Document
 from langchain_core.language_models import BaseChatModel
-
 from ragstack_knowledge_graph.schema_inference import KnowledgeSchemaInferer
 
 MARIE_CURIE_SOURCE = """
diff --git a/libs/knowledge-store/notebooks/astra_support.ipynb b/libs/knowledge-store/notebooks/astra_support.ipynb
@@ -40,6 +40,9 @@
    "metadata": {},
    "outputs": [],
    "source": [
+    "import requests\n",
+    "from bs4 import BeautifulSoup\n",
+    "\n",
     "# Use sitemaps to crawl the content\n",
     "SITEMAPS = [\n",
     "    \"https://docs.datastax.com/en/sitemap-astra-db-vector.xml\",\n",
@@ -54,9 +57,6 @@
     "\n",
     "SITE_PREFIX = \"astra\"\n",
     "\n",
-    "from bs4 import BeautifulSoup\n",
-    "import requests\n",
-    "\n",
     "\n",
     "def load_pages(sitemap_url):\n",
     "    r = requests.get(\n",
@@ -99,14 +99,14 @@
    "metadata": {},
    "outputs": [],
    "source": [
+    "from typing import AsyncIterator, Iterable\n",
+    "\n",
     "from langchain_community.document_loaders import AsyncHtmlLoader\n",
-    "from bs4 import BeautifulSoup\n",
     "from langchain_core.documents import Document\n",
-    "from typing import AsyncIterator, Iterable\n",
-    "from ragstack_knowledge_store.graph_store import CONTENT_ID\n",
     "from markdownify import MarkdownConverter\n",
+    "from ragstack_knowledge_store.graph_store import CONTENT_ID\n",
+    "from ragstack_langchain.graph_store.extractors import HtmlInput, HtmlLinkExtractor\n",
     "from ragstack_langchain.graph_store.links import add_links\n",
-    "from ragstack_langchain.graph_store.extractors import HtmlLinkExtractor, HtmlInput\n",
     "\n",
     "markdown_converter = MarkdownConverter(heading_style=\"ATX\")\n",
     "html_link_extractor = HtmlLinkExtractor()\n",
@@ -228,7 +228,7 @@
     "    import cassio\n",
     "\n",
     "    cassio.init(auto=True)\n",
-    "    from cassio.config import check_resolve_session, check_resolve_keyspace\n",
+    "    from cassio.config import check_resolve_keyspace, check_resolve_session\n",
     "\n",
     "    session = check_resolve_session()\n",
     "    keyspace = check_resolve_keyspace()\n",
@@ -252,7 +252,9 @@
     "cassio.init(auto=True)\n",
     "embeddings = OpenAIEmbeddings()\n",
     "graph_store = CassandraGraphStore(\n",
-    "    embeddings, node_table=f\"{SITE_PREFIX}_nodes\", targets_table=f\"{SITE_PREFIX}_targets\"\n",
+    "    embeddings,\n",
+    "    node_table=f\"{SITE_PREFIX}_nodes\",\n",
+    "    targets_table=f\"{SITE_PREFIX}_targets\",\n",
     ")"
    ]
   },
@@ -304,10 +306,10 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "from langchain_openai import ChatOpenAI\n",
+    "from langchain_core.output_parsers import StrOutputParser\n",
     "from langchain_core.prompts import ChatPromptTemplate\n",
     "from langchain_core.runnables import RunnablePassthrough\n",
-    "from langchain_core.output_parsers import StrOutputParser\n",
+    "from langchain_openai import ChatOpenAI\n",
     "\n",
     "llm = ChatOpenAI(model=\"gpt-4o\")\n",
     "\n",
@@ -348,7 +350,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "from IPython.display import display, Markdown\n",
+    "from IPython.display import Markdown, display\n",
     "\n",
     "\n",
     "# Helper method to render markdown in responses to a chain.\n",
diff --git a/libs/knowledge-store/notebooks/pdf_keybert.ipynb b/libs/knowledge-store/notebooks/pdf_keybert.ipynb
@@ -72,7 +72,9 @@
     "\n",
     "os.environ[\"OPENAI_API_KEY\"] = getpass.getpass(\"Enter OpenAI API Key: \")\n",
     "os.environ[\"ASTRA_DB_DATABASE_ID\"] = input(\"Enter Astra DB Database ID: \")\n",
-    "os.environ[\"ASTRA_DB_APPLICATION_TOKEN\"] = getpass.getpass(\"Enter Astra DB Application Token: \")\n",
+    "os.environ[\"ASTRA_DB_APPLICATION_TOKEN\"] = getpass.getpass(\n",
+    "    \"Enter Astra DB Application Token: \"\n",
+    ")\n",
     "\n",
     "keyspace = input(\"Enter Astra DB Keyspace (Empty for default): \")\n",
     "if keyspace:\n",
@@ -91,6 +93,7 @@
     "# See `env.template` for an example of what you should have there.\n",
     "%pip install python-dotenv\n",
     "import dotenv\n",
+    "\n",
     "dotenv.load_dotenv()"
    ]
   },
@@ -129,6 +132,7 @@
    "source": [
     "# Initialize cassandra connection from environment variables).\n",
     "import cassio\n",
+    "\n",
     "cassio.init(auto=True)"
    ]
   },
@@ -139,8 +143,9 @@
    "outputs": [],
    "source": [
     "# Create graph store.\n",
-    "from ragstack_knowledge_store import KnowledgeStore\n",
-    "graph_store = GraphStore(embeddings)"
+    "from ragstack_langchain.graph_store import CassandraGraphStore\n",
+    "\n",
+    "graph_store = CassandraGraphStore(embeddings)"
    ]
   },
   {
@@ -191,10 +196,11 @@
     "from keybert import KeyBERT\n",
     "\n",
     "kw_model = KeyBERT()\n",
-    "keywords = kw_model.extract_keywords([doc.page_content for doc in pages],\n",
-    "                                     stop_words='english')\n",
+    "keywords = kw_model.extract_keywords(\n",
+    "    [doc.page_content for doc in pages], stop_words=\"english\"\n",
+    ")\n",
     "\n",
-    "for (doc, kws) in zip(pages, keywords):\n",
+    "for doc, kws in zip(pages, keywords):\n",
     "    # Consider only taking keywords within a certain distance?\n",
     "    doc.metadata[\"keywords\"] = [kw for (kw, _) in kws]\n",
     "pages[0]"
@@ -239,9 +245,9 @@
    "outputs": [],
    "source": [
     "# Retrieve and generate using the relevant snippets of the blog.\n",
-    "from langchain_core.runnables import RunnablePassthrough\n",
     "from langchain_core.output_parsers import StrOutputParser\n",
     "from langchain_core.prompts import ChatPromptTemplate\n",
+    "from langchain_core.runnables import RunnablePassthrough\n",
     "\n",
     "retriever0 = graph_store.as_retriever(depth=0)\n",
     "retriever1 = graph_store.as_retriever(depth=1)\n",
@@ -253,6 +259,7 @@
     "\"\"\"\n",
     "prompt = ChatPromptTemplate.from_template(template)\n",
     "\n",
+    "\n",
     "def format_docs(docs):\n",
     "    return \"\\n\\n\".join(doc.page_content for doc in docs)\n",
     "\n",
diff --git a/libs/knowledge-store/pyproject.toml b/libs/knowledge-store/pyproject.toml
@@ -33,30 +33,6 @@ numpy = [
 requires = ["poetry-core"]
 build-backend = "poetry.core.masonry.api"
 
-
-[tool.ruff]
-line-length = 98
-
-[tool.ruff.lint]
-# Enable Pyflakes `E` and `F` codes by default.
-select = [
-    # Pyflakes
-    "F",
-    # Pycodestyle
-    "E",
-    "W",
-    # isort
-    "I001",
-]
-ignore = []
-
-# Allow unused variables when underscore-prefixed.
-dummy-variable-rgx = "^(_+|(_+[a-zA-Z0-9_]*[a-zA-Z0-9]+?))$"
-
-[tool.ruff.lint.mccabe]
-# Unlike Flake8, default to a complexity level of 10.
-max-complexity = 10
-
 [tool.mypy]
 strict = true
 warn_unreachable = true
diff --git a/libs/knowledge-store/ragstack_knowledge_store/graph_store.py b/libs/knowledge-store/ragstack_knowledge_store/graph_store.py
