Add boolean type (eyurtsev#165)

This PR adds a boolean type to Kor.

docs/source/api.rst

@@ -33,8 +33,9 @@ representation automatically.
 .. autosummary::
 
  Object
- Number
  Text
+ Number
+ Bool
  Selection
  Option
 

kor/__init__.py

@@ -6,7 +6,7 @@
  create_extraction_chain,
  extract_from_documents,
 )
-from .nodes import Number, Object, Option, Selection, Text
+from .nodes import Bool, Number, Object, Option, Selection, Text
 from .type_descriptors import (
  BulletPointDescriptor,
  TypeDescriptor,
@@ -22,6 +22,7 @@
  "Extraction",
  "from_pydantic",
  "JSONEncoder",
+ "Bool",
  "Number",
  "Object",
  "Option",

kor/adapters.py

@@ -15,7 +15,7 @@
 
 from pydantic import BaseModel
 
-from .nodes import ExtractionSchemaNode, Number, Object, Option, Selection, Text
+from .nodes import Bool, ExtractionSchemaNode, Number, Object, Option, Selection, Text
 from .validators import PydanticValidator, Validator
 
 # Not going to support dicts or lists since that requires recursive checks.
@@ -86,7 +86,7 @@ def _translate_pydantic_to_kor(
  name=field_name,
  )
  elif issubclass(type_, bool):
- attribute = Text(
+ attribute = Bool(
  id=field_name,
  examples=field_examples,
  description=field_description,

kor/documents/html.py

@@ -3,8 +3,6 @@
 import re
 from typing import Tuple
 
-import markdownify
-from bs4 import BeautifulSoup
 from langchain.schema import Document
 
 from kor.documents.typedefs import AbstractDocumentProcessor
@@ -16,6 +14,13 @@
 
 def _get_mini_html(html: str, *, tags_to_remove: Tuple[str, ...] = tuple()) -> str:
  """Clean up HTML tags."""
+ try:
+ from bs4 import BeautifulSoup
+ except ImportError:
+ raise ImportError(
+ "Please install BeautifulSoup to use the HTML document processor. "
+ "You can do so by running `pip install beautifulsoup4`."
+ )
  # Parse the HTML document using BeautifulSoup
  soup = BeautifulSoup(html, "html.parser")
 
@@ -34,6 +39,14 @@ def _get_mini_html(html: str, *, tags_to_remove: Tuple[str, ...] = tuple()) -> s
 
 def _clean_html(html: str, *, tags_to_remove: Tuple[str, ...] = tuple()) -> str:
  """Clean up HTML and convert to markdown using markdownify."""
+ try:
+ import markdownify
+ except ImportError:
+ raise ImportError(
+ "Please install markdownify to use the HTML document processor. "
+ "You can do so by running `pip install markdownify`."
+ )
+
  html = _get_mini_html(html, tags_to_remove=tags_to_remove)
  md = markdownify.markdownify(html)
  return CONSECUTIVE_NEW_LINES.sub("\n\n", md).strip()

kor/extraction/api.py

@@ -9,7 +9,7 @@
 try: # Handle breaking change in langchain
  from langchain.base_language import BaseLanguageModel
 except ImportError:
- from langchain.schema import BaseLanguageModel
+ from langchain.schema import BaseLanguageModel # type: ignore
 
 from kor.encoders import Encoder, InputFormatter, initialize_encoder
 from kor.extraction.typedefs import DocumentExtraction, Extraction

kor/nodes.py

@@ -1,4 +1,6 @@
 """Definitions of input elements."""
+from __future__ import annotations
+
 import abc
 import copy
 import re
@@ -31,30 +33,34 @@
 class AbstractVisitor(Generic[T], abc.ABC):
  """An abstract visitor."""
 
- def visit_text(self, node: "Text", **kwargs: Any) -> T:
+ def visit_text(self, node: Text, **kwargs: Any) -> T:
  """Visit text node."""
  return self.visit_default(node, **kwargs)
 
- def visit_number(self, node: "Number", **kwargs: Any) -> T:
+ def visit_number(self, node: Number, **kwargs: Any) -> T:
  """Visit text node."""
  return self.visit_default(node, **kwargs)
 
- def visit_object(self, node: "Object", **kwargs: Any) -> T:
+ def visit_object(self, node: Object, **kwargs: Any) -> T:
  """Visit object node."""
  return self.visit_default(node, **kwargs)
 
- def visit_selection(self, node: "Selection", **kwargs: Any) -> T:
+ def visit_selection(self, node: Selection, **kwargs: Any) -> T:
  """Visit selection node."""
  return self.visit_default(node, **kwargs)
 
- def visit_option(self, node: "Option", **kwargs: Any) -> T:
+ def visit_option(self, node: Option, **kwargs: Any) -> T:
  """Visit option node."""
  return self.visit_default(node, **kwargs)
 
- def visit_default(self, node: "AbstractSchemaNode", **kwargs: Any) -> T:
+ def visit_default(self, node: AbstractSchemaNode, **kwargs: Any) -> T:
  """Default node implementation."""
  raise NotImplementedError()
 
+ def visit_bool(self, node: Bool, **kwargs: Any) -> T:
+ """Visit bool node."""
+ return self.visit_default(node, **kwargs)
+
 
 class AbstractSchemaNode(BaseModel):
  """Abstract schema node.
@@ -140,6 +146,14 @@ def accept(self, visitor: AbstractVisitor[T], **kwargs: Any) -> T:
  return visitor.visit_text(self, **kwargs)
 
 
+class Bool(ExtractionSchemaNode):
+ """Built-in bool input."""
+
+ def accept(self, visitor: AbstractVisitor[T], **kwargs: Any) -> T:
+ """Accept a visitor."""
+ return visitor.visit_bool(self, **kwargs)
+
+
 class Option(AbstractSchemaNode):
  """Built-in option input must be part of a selection input."""
 
@@ -219,7 +233,7 @@ class Object(AbstractSchemaNode):
 
  """
 
- attributes: Sequence[Union[ExtractionSchemaNode, Selection, "Object"]]
+ attributes: Sequence[Union[ExtractionSchemaNode, Selection, Object]]
 
  examples: Sequence[
  Tuple[

kor/type_descriptors.py

@@ -13,6 +13,7 @@
 from kor.nodes import (
  AbstractSchemaNode,
  AbstractVisitor,
+ Bool,
  Number,
  Object,
  Selection,
@@ -76,6 +77,8 @@ def visit_default(self, node: "AbstractSchemaNode", **kwargs: Any) -> List[str]:
  finalized_type = "string"
  elif isinstance(node, Number):
  finalized_type = "number"
+ elif isinstance(node, Bool):
+ finalized_type = "boolean"
  else:
  raise NotImplementedError()