Add type checking of knowledge-graph (#583)

Author: cbornet

.github/workflows/ci-unit-tests.yml

@@ -52,8 +52,8 @@ jobs:
           cd docker/examples/local-llm
           sudo docker build -t local-llm .
 
-  unit-tests:
-    name: Unit Tests (Python ${{ matrix.python-version }})
+  lint:
+    name: Lint (Python ${{ matrix.python-version }})
     needs: ["preconditions"]
     runs-on: ubuntu-latest
     strategy:
@@ -73,9 +73,43 @@ jobs:
           python-version: "${{ matrix.python-version }}"
 
       - name: Run lint
-        if: ${{ matrix.python-version != '3.12' }}
         uses: ./.github/actions/lint
 
+  type-check:
+    name: Type Check
+    needs: ["preconditions"]
+    runs-on: ubuntu-latest
+    steps:
+      - name: Check out the repo
+        uses: actions/checkout@v4
+
+      - name: "Setup: Python 3.11"
+        uses: ./.github/actions/setup-python
+
+      - name: "Type check (knowledge-graph)"
+        run: tox -e type -c libs/knowledge-graph && rm -rf libs/knowledge-graph/.tox
+
+
+  unit-tests:
+    name: Unit Tests (Python ${{ matrix.python-version }})
+    needs: ["preconditions"]
+    runs-on: ubuntu-latest
+    strategy:
+      matrix:
+        python-version:
+          - "3.12"
+          - "3.11"
+          - "3.10"
+          - "3.9"
+    steps:
+      - name: Check out the repo
+        uses: actions/checkout@v4
+
+      - name: "Setup: Python ${{ matrix.python-version }}"
+        uses: ./.github/actions/setup-python
+        with:
+          python-version: "${{ matrix.python-version }}"
+
       - name: "Unit tests (root)"
         # yamllint disable-line rule:line-length
         if: ${{ needs.preconditions.outputs.libs_langchain == 'true' || needs.preconditions.outputs.libs_colbert == 'true' || needs.preconditions.outputs.libs_llamaindex == 'true' }}

libs/knowledge-graph/pyproject.toml

@@ -34,18 +34,30 @@ pytest = "^8.1.1"
 pytest-asyncio = "^0.23.6"
 pytest-dotenv = "^0.5.2"
 pytest-rerunfailures = "^14.0"
-setuptools = "^70.0.0"
+mypy = "^1.10.1"
+types-pyyaml = "^6.0.1"
+pydantic = "<2" # for compatibility between LangChain and pydantic-yaml type checking
 
 [build-system]
 requires = ["poetry-core"]
 build-backend = "poetry.core.masonry.api"
 
 [tool.mypy]
-strict = true
-warn_unreachable = true
-pretty = true
-show_column_numbers = true
+disallow_any_generics = true
+disallow_incomplete_defs = true
+disallow_untyped_calls = true
+disallow_untyped_decorators = true
+disallow_untyped_defs = true
+follow_imports = "normal"
+ignore_missing_imports = true
+no_implicit_reexport = true
+show_error_codes = true
 show_error_context = true
+strict_equality = true
+strict_optional = true
+warn_redundant_casts = true
+warn_return_any = true
+warn_unused_ignores = true
 
 [tool.pytest.ini_options]
 testpaths = ["tests"]

libs/knowledge-graph/ragstack_knowledge_graph/cassandra_graph_store.py

@@ -58,7 +58,9 @@ def add_graph_documents(
 
     # TODO: should this include the types of each node?
     @override
-    def query(self, query: str, params: dict = {}) -> List[Dict[str, Any]]:  # noqa: B006
+    def query(
+        self, query: str, params: Optional[Dict[str, Any]] = None
+    ) -> List[Dict[str, Any]]:
         raise ValueError("Querying Cassandra should use `as_runnable`.")
 
     @override
@@ -76,7 +78,9 @@ def get_structured_schema(self) -> Dict[str, Any]:
     def refresh_schema(self) -> None:
         raise NotImplementedError
 
-    def as_runnable(self, steps: int = 3, edge_filters: Sequence[str] = ()) -> Runnable:
+    def as_runnable(
+        self, steps: int = 3, edge_filters: Sequence[str] = ()
+    ) -> Runnable[Union[Node, Sequence[Node]], Iterable[Relation]]:
         """Convert to a runnable.
 
         Returns a runnable that retrieves the sub-graph near the input entity or

libs/knowledge-graph/ragstack_knowledge_graph/extraction.py

@@ -1,8 +1,9 @@
-from typing import Dict, List, Sequence, Union, cast
+from typing import TYPE_CHECKING, Any, Dict, List, Sequence, Union, cast
 
 from langchain_community.graphs.graph_document import GraphDocument
 from langchain_core.documents import Document
 from langchain_core.language_models.chat_models import BaseChatModel
+from langchain_core.messages import SystemMessage
 from langchain_core.prompts import (
     ChatPromptTemplate,
     HumanMessagePromptTemplate,
@@ -23,6 +24,9 @@
 )
 from .templates import load_template
 
+if TYPE_CHECKING:
+    from langchain_core.prompts.chat import MessageLikeRepresentation
+
 
 def _format_example(idx: int, example: Example) -> str:
     from pydantic_yaml import to_yaml_str
@@ -43,7 +47,7 @@ def __init__(
         self._validator = KnowledgeSchemaValidator(schema)
         self.strict = strict
 
-        messages = [
+        messages: List[MessageLikeRepresentation] = [
             SystemMessagePromptTemplate(
                 prompt=load_template(
                     "extraction.md", knowledge_schema_yaml=schema.to_yaml_str()
@@ -52,22 +56,24 @@ def __init__(
         ]
 
         if examples:
-            formatted = "\n\n".join(map(_format_example, examples))
-            messages.append(SystemMessagePromptTemplate(prompt=formatted))
+            formatted = "\n\n".join(
+                [_format_example(i, example) for i, example in enumerate(examples)]
+            )
+            messages.append(SystemMessage(content=formatted))
 
         messages.append(HumanMessagePromptTemplate.from_template("Input: {input}"))
 
         prompt = ChatPromptTemplate.from_messages(messages)
-        schema = create_simple_model(
+        model_schema = create_simple_model(
             node_labels=[node.type for node in schema.nodes],
             rel_types=list({r.edge_type for r in schema.relationships}),
         )
         # TODO: Use "full" output so we can detect parsing errors?
-        structured_llm = llm.with_structured_output(schema)
+        structured_llm = llm.with_structured_output(model_schema)
         self._chain = prompt | structured_llm
 
     def _process_response(
-        self, document: Document, response: Union[Dict, BaseModel]
+        self, document: Document, response: Union[Dict[str, Any], BaseModel]
     ) -> GraphDocument:
         raw_graph = cast(_Graph, response)
         nodes = (
@@ -81,14 +87,14 @@ def _process_response(
             else []
         )
 
-        document = GraphDocument(
+        graph_document = GraphDocument(
             nodes=nodes, relationships=relationships, source=document
         )
 
         if self.strict:
-            self._validator.validate_graph_document(document)
+            self._validator.validate_graph_document(graph_document)
 
-        return document
+        return graph_document
 
     def extract(self, documents: List[Document]) -> List[GraphDocument]:
         """Extract knowledge graphs from a list of documents."""

libs/knowledge-graph/ragstack_knowledge_graph/knowledge_graph.py

@@ -19,7 +19,7 @@ def _deserialize_md_dict(md_string: str) -> Dict[str, Any]:
     return cast(Dict[str, Any], json.loads(md_string))
 
 
-def _parse_node(row) -> Node:
+def _parse_node(row: Any) -> Node:
     return Node(
         name=row.name,
         type=row.type,
@@ -106,7 +106,7 @@ def __init__(
             """
         )
 
-    def _apply_schema(self):
+    def _apply_schema(self) -> None:
         # Partition by `name` and cluster by `type`.
         # Each `(name, type)` pair is a unique node.
         # We can enumerate all `type` values for a given `name` to identify ambiguous
@@ -152,11 +152,13 @@ def _apply_schema(self):
             """
         )
 
-    def _send_query_nearest_node(self, node: str, k: int = 1) -> ResponseFuture:
+    def _send_query_nearest_node(
+        self, embeddings: Embeddings, node: str, k: int = 1
+    ) -> ResponseFuture:
         return self._session.execute_async(
             self._query_nodes_by_embedding,
             (
-                self._text_embeddings.embed_query(node),
+                embeddings.embed_query(node),
                 k,
             ),
         )
@@ -173,13 +175,11 @@ def query_nearest_nodes(self, nodes: Iterable[str], k: int = 1) -> Iterable[Node
             raise ValueError("Unable to query for nearest nodes without embeddings")
 
         node_futures: Iterable[ResponseFuture] = [
-            self._send_query_nearest_node(n, k) for n in nodes
+            self._send_query_nearest_node(self._text_embeddings, n, k) for n in nodes
         ]
-
-        nodes = {
+        return {
             _parse_node(n) for node_future in node_futures for n in node_future.result()
         }
-        return list(nodes)
 
     # TODO: Introduce `ainsert` for async insertions.
     def insert(
@@ -253,9 +253,11 @@ def subgraph(
             for n in nodes
         ]
 
-        nodes = [_parse_node(n) for future in node_futures for n in future.result()]
+        graph_nodes = [
+            _parse_node(n) for future in node_futures for n in future.result()
+        ]
 
-        return nodes, edges
+        return graph_nodes, edges
 
     def traverse(
         self,

libs/knowledge-graph/ragstack_knowledge_graph/knowledge_schema.py

@@ -1,5 +1,5 @@
 from pathlib import Path
-from typing import Dict, List, Sequence, Union
+from typing import Dict, List, Self, Sequence, Union
 
 from langchain_community.graphs.graph_document import GraphDocument
 from langchain_core.pydantic_v1 import BaseModel
@@ -66,7 +66,7 @@ class KnowledgeSchema(BaseModel):
     """Allowed relationships for the knowledge schema."""
 
     @classmethod
-    def from_file(cls, path: Union[str, Path]) -> "KnowledgeSchema":
+    def from_file(cls, path: Union[str, Path]) -> Self:
         """Load a KnowledgeSchema from a JSON or YAML file.
 
         Args:
@@ -98,27 +98,27 @@ def __init__(self, schema: KnowledgeSchema) -> None:
             # TODO: Validate the relationship.
             # source/target type should exist in nodes, edge_type should exist in edges
 
-    def validate_graph_document(self, document: GraphDocument):
+    def validate_graph_document(self, document: GraphDocument) -> None:
         """Validate a graph document against the schema."""
         e = ValueError("Invalid graph document for schema")
         for node_type in {node.type for node in document.nodes}:
             if node_type not in self._nodes:
                 e.add_note(f"No node type '{node_type}")
         for r in document.relationships:
-            relationships = self._relationships.get(r.edge_type, None)
+            relationships = self._relationships.get(r.type, None)
             if relationships is None:
-                e.add_note(f"No edge type '{r.edge_type}")
+                e.add_note(f"No edge type '{r.type}")
             else:
                 relationship = next(
                     candidate
                     for candidate in relationships
-                    if r.source_type in candidate.source_types
-                    if r.target_type in candidate.target_types
+                    if r.source.type in candidate.source_types
+                    if r.target.type in candidate.target_types
                 )
                 if relationship is None:
                     e.add_note(
                         "No relationship allows "
-                        f"({r.source_id} -> {r.type} -> {r.target.type})"
+                        f"({r.source.id} -> {r.type} -> {r.target.type})"
                     )
 
         if e.__notes__:

libs/knowledge-graph/ragstack_knowledge_graph/render.py

@@ -1,4 +1,4 @@
-from typing import Iterable, Union
+from typing import Dict, Iterable, Tuple, Union
 
 import graphviz
 from langchain_community.graphs.graph_document import GraphDocument, Node
@@ -12,7 +12,7 @@ def _node_label(node: Node) -> str:
 
 def print_graph_documents(
     graph_documents: Union[GraphDocument, Iterable[GraphDocument]],
-):
+) -> None:
     """Prints the relationships in the graph documents."""
     if isinstance(graph_documents, GraphDocument):
         graph_documents = [graph_documents]
@@ -29,13 +29,13 @@ def render_graph_documents(
 ) -> graphviz.Digraph:
     """Renders the relationships in the graph documents."""
     if isinstance(graph_documents, GraphDocument):
-        graph_documents = [GraphDocument]
+        graph_documents = [graph_documents]
 
     dot = graphviz.Digraph()
 
-    nodes = {}
+    nodes: Dict[Tuple[Union[str, int], str], str] = {}
 
-    def _node_id(node: Node) -> int:
+    def _node_id(node: Node) -> str:
         node_key = (node.id, node.type)
         if node_id := nodes.get(node_key):
             return node_id

libs/knowledge-graph/ragstack_knowledge_graph/runnables.py

@@ -1,4 +1,4 @@
-from typing import List, Optional
+from typing import Any, Dict, List, Optional
 
 from langchain_core.language_models import BaseChatModel
 from langchain_core.output_parsers import JsonOutputParser
@@ -26,7 +26,7 @@ def extract_entities(
     llm: BaseChatModel,
     keyword_extraction_prompt: str = QUERY_ENTITY_EXTRACT_PROMPT,
     node_types: Optional[List[str]] = None,
-) -> Runnable:
+) -> Runnable[Dict[str, Any], List[Node]]:
     """Return a keyword-extraction runnable.
 
     This will expect a dictionary containing the `"question"` to extract keywords from.

libs/knowledge-graph/ragstack_knowledge_graph/templates.py

@@ -1,5 +1,5 @@
 from os import path
-from typing import Callable, Union
+from typing import Callable, Union, cast
 
 from langchain_core.prompts import PromptTemplate
 
@@ -12,5 +12,5 @@ def load_template(
     """Load a template from a file."""
     template = PromptTemplate.from_file(path.join(TEMPLATE_PATH, filename))
     if kwargs:
-        template = template.partial(**kwargs)
+        template = cast(PromptTemplate, template.partial(**kwargs))
     return template