Skip to content

API Reference

Classes

Base extractor

BaseExtractior

A class to handle any extraction using OpenAI's GPT models.

Source code in code/NLI_base_extractor.py 
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
class BaseExtractior:
    """
    A class to handle any extraction using OpenAI's GPT models.
    """

    def __init__(self, openai_client: OpenAI | AsyncOpenAI) -> None:
        """
        Initializes the BaseExtractior with an OpenAI client.

        Args:
            openai_client (OpenAI | AsyncOpenAI): An instance of OpenAI or AsyncOpenAI client to interact with the API.
        """
        self.openai_client = openai_client
        self.is_async = False
        if isinstance(self.openai_client, AsyncOpenAI):
            self.is_async = True

    def _get_completion_result(
        self, model, messages: list[dict], temperature: float
    ) -> str:

        completion = self.openai_client.chat.completions.create(
            model=model,
            messages=messages,
            temperature=temperature,
        )

        return completion.choices[0].message.content

    def _get_completion_parsed_result(
        self, model, messages: list[dict], response_format, temperature: float, key: str
    ):

        completion = self.openai_client.beta.chat.completions.parse(
            model=model,
            messages=messages,
            response_format=response_format,
            temperature=temperature,
        )

        return json.loads(completion.choices[0].message.content)[key]

__init__(openai_client)

Initializes the BaseExtractior with an OpenAI client.

Parameters:

Name Type Description Default
openai_client OpenAI | AsyncOpenAI

An instance of OpenAI or AsyncOpenAI client to interact with the API.

 required
Source code in code/NLI_base_extractor.py 
11
12
13
14
15
16
17
18
19
20
21
def __init__(self, openai_client: OpenAI | AsyncOpenAI) -> None:
    """
    Initializes the BaseExtractior with an OpenAI client.

    Args:
        openai_client (OpenAI | AsyncOpenAI): An instance of OpenAI or AsyncOpenAI client to interact with the API.
    """
    self.openai_client = openai_client
    self.is_async = False
    if isinstance(self.openai_client, AsyncOpenAI):
        self.is_async = True

Node extractor

NodeExtractor

Bases: BaseExtractior

A class to handle node extraction from a given text description using OpenAI's GPT models.

Source code in code/NLI_node_extraction.py 
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
class NodeExtractor(BaseExtractior):
    """
    A class to handle node extraction from a given text description using OpenAI's GPT models.
    """

    def __init__(self, openai_client: OpenAI | AsyncOpenAI) -> None:
        """
        Initializes the NodeExtractor with an OpenAI client.

        Args:
            openai_client (OpenAI | AsyncOpenAI): An instance of OpenAI or AsyncOpenAI client to interact with the API.
        """
        super(NodeExtractor, self).__init__(openai_client)

    def extract_nodes_gpt(
        self, description: str, gpt_model: str = "gpt-4o-mini", temperature: float = 0.0
    ) -> list[str]:
        """
        Extracts nodes from a given description using a specified GPT model.

        Args:
            description (str): The text description from which nodes need to be extracted.
            gpt_model (str, optional): The GPT model to use for extraction. Defaults to 'gpt-4o-mini'.
            temperature (float, optional): The randomness of the model's output. Defaults to 0.0.

        Returns:
            list[str]: A list of extracted nodes from the description.
        """

        list_of_nodes = self._get_completion_parsed_result(
            model=gpt_model,
            messages=self._get_messages_for_node_extraction(description),
            response_format=Nodes,
            temperature=temperature,
            key="list_of_nodes",
        )

        return list_of_nodes

    @staticmethod
    def _get_messages_for_node_extraction(description: str) -> tuple[dict, dict]:
        """
        Prepares the message payload for node extraction using the GPT model.

        Args:
            description (str): The description text to be processed.

        Returns:
            tuple[dict, dict]: A tuple containing system and user messages formatted for the API request.
        """
        messages = (
            {"role": "system", "content": node_extraction_sys_message},
            {
                "role": "user",
                "content": node_extraction_str_template.format(description=description),
            },
        )

        return messages

__init__(openai_client)

Initializes the NodeExtractor with an OpenAI client.

Parameters:

Name Type Description Default
openai_client OpenAI | AsyncOpenAI

An instance of OpenAI or AsyncOpenAI client to interact with the API.

 required
Source code in code/NLI_node_extraction.py 
17
18
19
20
21
22
23
24
def __init__(self, openai_client: OpenAI | AsyncOpenAI) -> None:
    """
    Initializes the NodeExtractor with an OpenAI client.

    Args:
        openai_client (OpenAI | AsyncOpenAI): An instance of OpenAI or AsyncOpenAI client to interact with the API.
    """
    super(NodeExtractor, self).__init__(openai_client)

extract_nodes_gpt(description, gpt_model='gpt-4o-mini', temperature=0.0)

Extracts nodes from a given description using a specified GPT model.

Parameters:

Name Type Description Default
description str

The text description from which nodes need to be extracted.

 required
gpt_model str

The GPT model to use for extraction. Defaults to 'gpt-4o-mini'.

 'gpt-4o-mini'
temperature float

The randomness of the model's output. Defaults to 0.0.

 0.0

Returns:

Type Description
list[str]

list[str]: A list of extracted nodes from the description.
Source code in code/NLI_node_extraction.py 
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
def extract_nodes_gpt(
    self, description: str, gpt_model: str = "gpt-4o-mini", temperature: float = 0.0
) -> list[str]:
    """
    Extracts nodes from a given description using a specified GPT model.

    Args:
        description (str): The text description from which nodes need to be extracted.
        gpt_model (str, optional): The GPT model to use for extraction. Defaults to 'gpt-4o-mini'.
        temperature (float, optional): The randomness of the model's output. Defaults to 0.0.

    Returns:
        list[str]: A list of extracted nodes from the description.
    """

    list_of_nodes = self._get_completion_parsed_result(
        model=gpt_model,
        messages=self._get_messages_for_node_extraction(description),
        response_format=Nodes,
        temperature=temperature,
        key="list_of_nodes",
    )

    return list_of_nodes

Edge extrctor

Code for vertices extraction.

This module provides functionality to extract edges between nodes in a graph using OpenAI's GPT models. It includes methods to determine the existence and direction of edges based on a textual description.

The module contains the following functions:

  • _extract_one_edge_gpt: Determines the existence and direction of an edge between a pair of nodes.
  • extract_all_edges: Extracts all possible edges from a set of nodes based on a given description.

ArrowEnum

Bases: str, Enum

Enumeration for arrow direction types in a graph.

Source code in code/NLI_extract_edges.py 
18
19
20
21
22
23
24
25
class ArrowEnum(str, Enum):
    """
    Enumeration for arrow direction types in a graph.
    """

    no = "no arrow"
    forward = "forward arrow"
    backward = "backward arrow"

ArrowType

Bases: BaseModel

Model representing the type of arrow direction for an edge.

Source code in code/NLI_extract_edges.py 
28
29
30
31
32
33
class ArrowType(BaseModel):
    """
    Model representing the type of arrow direction for an edge.
    """

    arrow_type: ArrowEnum

EdgeExtractor

Bases: BaseExtractior

Class to handle edge extraction from a description using OpenAI's GPT models.

Source code in code/NLI_extract_edges.py 
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
class EdgeExtractor(BaseExtractior):
    """
    Class to handle edge extraction from a description using OpenAI's GPT models.
    """

    def __init__(self, openai_client: OpenAI | AsyncOpenAI) -> None:
        """
        Initializes the NodeExtraction with an OpenAI client.

        Args:
            openai_client (OpenAI | AsyncOpenAI): An instance of OpenAI or AsyncOpenAI client to interact with the API.
        """
        super(EdgeExtractor, self).__init__(openai_client)

    @staticmethod
    def _get_messages_for_edge_direction(
        description: str,
        set_of_nodes: list[str],
        pair_of_nodes: tuple[str, str],
    ) -> tuple[dict, dict]:
        """
        Prepares the message payload for edge direction identification using the GPT model.

        Args:
            description (str): The description text to be processed.
            set_of_nodes (list[str]): The nodes of the graph.
            pair_of_nodes (tuple[str, str]): Pair of nodes to decide about the existance and the direction of the edge.

        Returns:
            tuple[dict, dict]: A tuple containing system and user messages formatted for the API request.
        """

        messages = (
            {"role": "system", "content": node_extraction_sys_message},
            {
                "role": "user",
                "content": edge_extraction_str_template.format(
                    description=description,
                    set_of_nodes=set_of_nodes,
                    pair_of_nodes=pair_of_nodes,
                ),
            },
        )

        return messages

    def _extract_one_edge_gpt(
        self,
        description: str,
        set_of_nodes: list[str],
        pair_of_nodes: tuple[str, str],
        gpt_model: str = "gpt-4o-mini",
        temperature: float = 0,
    ) -> tuple[str | None, str | None]:
        """
        Determines the existence and direction of an edge between a pair of nodes using a GPT model.

        Args:
            description (str): The description text to be processed.
            set_of_nodes (list[str]): The nodes of the graph.
            pair_of_nodes (tuple[str, str]): A pair of nodes to check for an edge.
            gpt_model (str, optional): The GPT model to use for extraction. Defaults to 'gpt-4o-mini'.
            temperature (float, optional): The randomness of the model's output. Defaults to 0.

        Returns:
            tuple[str | None, str | None]: either (None, None) if no edge exists, or a tuple representing the edge with identified direction.
        """

        arrow_type = self._get_completion_parsed_result(
            model=gpt_model,
            messages=self._get_messages_for_edge_direction(
                description,
                f"[{', '.join(set_of_nodes)}]",
                f"[{', '.join(pair_of_nodes)}]",
            ),
            response_format=ArrowType,
            temperature=temperature,
            key="arrow_type",
        )

        if "forward" in arrow_type.lower():
            return pair_of_nodes
        if "backward" in arrow_type.lower():
            return pair_of_nodes[::-1]

        return (None, None)

    def extract_all_edges(
        self,
        description: str,
        set_of_nodes: list[str],
        gpt_model: str = "gpt-4o-mini",
        temperature: float = 0,
        verbose=False,
    ) -> list[tuple[str | None, str | None]]:
        """
        Extracts all possible edges from a set of nodes based on a given description.

        Args:
            description (str): The description text to be processed.
            set_of_nodes (list[str]): The nodes of the graph.
            gpt_model (str, optional): The GPT model to use for extraction. Defaults to 'gpt-4o-mini'.
            temperature (float, optional): The randomness of the model's output. Defaults to 0.
            verbose (bool, optional): If True, provides detailed logging. Defaults to False.

        Returns:
            list[tuple[str | None, str | None]]: A list of tuples representing edges with identified directions.
        """
        edge_list = []

        for i, node_a in enumerate(set_of_nodes):
            for node_b in set_of_nodes[i + 1 :]:
                if verbose:
                    print(f"{node_a} # {node_b}")
                edge = self._extract_one_edge_gpt(
                    description, set_of_nodes, (node_a, node_b), gpt_model, temperature
                )
                if edge[0] is not None:
                    edge_list.append(edge)

        return edge_list

__init__(openai_client)

Initializes the NodeExtraction with an OpenAI client.

Parameters:

Name Type Description Default
openai_client OpenAI | AsyncOpenAI

An instance of OpenAI or AsyncOpenAI client to interact with the API.

 required
Source code in code/NLI_extract_edges.py 
41
42
43
44
45
46
47
48
def __init__(self, openai_client: OpenAI | AsyncOpenAI) -> None:
    """
    Initializes the NodeExtraction with an OpenAI client.

    Args:
        openai_client (OpenAI | AsyncOpenAI): An instance of OpenAI or AsyncOpenAI client to interact with the API.
    """
    super(EdgeExtractor, self).__init__(openai_client)

extract_all_edges(description, set_of_nodes, gpt_model='gpt-4o-mini', temperature=0, verbose=False)

Extracts all possible edges from a set of nodes based on a given description.

Parameters:

Name Type Description Default
description str

The description text to be processed.

 required
set_of_nodes list[str]

The nodes of the graph.

 required
gpt_model str

The GPT model to use for extraction. Defaults to 'gpt-4o-mini'.

 'gpt-4o-mini'
temperature float

The randomness of the model's output. Defaults to 0.

 0
verbose bool

If True, provides detailed logging. Defaults to False.

 False

Returns:

Type Description
list[tuple[str | None, str | None]]

list[tuple[str | None, str | None]]: A list of tuples representing edges with identified directions.
Source code in code/NLI_extract_edges.py 
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
def extract_all_edges(
    self,
    description: str,
    set_of_nodes: list[str],
    gpt_model: str = "gpt-4o-mini",
    temperature: float = 0,
    verbose=False,
) -> list[tuple[str | None, str | None]]:
    """
    Extracts all possible edges from a set of nodes based on a given description.

    Args:
        description (str): The description text to be processed.
        set_of_nodes (list[str]): The nodes of the graph.
        gpt_model (str, optional): The GPT model to use for extraction. Defaults to 'gpt-4o-mini'.
        temperature (float, optional): The randomness of the model's output. Defaults to 0.
        verbose (bool, optional): If True, provides detailed logging. Defaults to False.

    Returns:
        list[tuple[str | None, str | None]]: A list of tuples representing edges with identified directions.
    """
    edge_list = []

    for i, node_a in enumerate(set_of_nodes):
        for node_b in set_of_nodes[i + 1 :]:
            if verbose:
                print(f"{node_a} # {node_b}")
            edge = self._extract_one_edge_gpt(
                description, set_of_nodes, (node_a, node_b), gpt_model, temperature
            )
            if edge[0] is not None:
                edge_list.append(edge)

    return edge_list

Node distribution extractor

This module contains the NodeDistributer class, which is designed to suggest distribution types for nodes based on textual descriptions using OpenAI's GPT models. It leverages natural language processing to infer whether a node's distribution should be binary, categorical, or continuous.

NodeDistributer

Bases: BaseExtractior

A class for distributing nodes based on text descriptions using OpenAI's GPT models.

Source code in code/NLI_suggest_node_distribution.py 
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
class NodeDistributer(BaseExtractior):
    """
    A class for distributing nodes based on text descriptions using OpenAI's GPT models.
    """

    def __init__(self, openai_client: OpenAI | AsyncOpenAI) -> None:
        """
        Initializes the NodeDistributer with an OpenAI client.

        Args:
            openai_client (OpenAI | AsyncOpenAI): An instance of OpenAI or AsyncOpenAI client to interact with the API.
        """

        super(NodeDistributer, self).__init__(openai_client)

    @staticmethod
    def _get_messages_for_vertex_distribution(
        description: str, node_name: str
    ) -> tuple[dict, dict]:
        """
        Prepares the message payload for node extraction using the GPT model.

        Args:
            description (str): The description text to be processed.
            node_name (str): Node name for which the distribution in inferred.

        Returns:
            tuple[dict, dict]: A tuple containing system and user messages formatted for the API request.
        """
        messages = (
            {"role": "system", "content": node_extraction_sys_message},
            {
                "role": "user",
                "content": node_distribution_str_template.format(
                    description=description, node_name=node_name
                ),
            },
        )

        return messages

    def _suggest_vertex_distribution(
        self,
        description: str,
        node_name: str,
        gpt_model: str = "gpt-4o-mini",
        temperature: float = 0,
    ) -> str:
        """
        Suggests the distribution type of a node based on its description.

        Args:
            description (str): The text description of the node.
            node_name (str): The name of the node for which to suggest a distribution.
            gpt_model (str, optional): The GPT model to use for suggestion. Defaults to 'gpt-4o-mini'.
            temperature (float, optional): The randomness of the model's output. Defaults to 0.

        Returns:
            str: The suggested distribution type, which can be 'binary', 'categorical', or 'continuous'.
        """
        distr_type = self._get_completion_result(
            model=gpt_model,
            messages=self._get_messages_for_vertex_distribution(description, node_name),
            temperature=temperature,
        )

        assert distr_type in [
            "binary",
            "categorical",
            "continuous",
        ], f"Wrong inferred distr type: {distr_type}"

        return distr_type

    def suggest_vertex_distributions(
        self,
        description: str,
        node_names: list[str],
        gpt_model: str = "gpt-4o-mini",
        temperature: float = 0,
    ) -> str:
        """
        Suggests distribution types for multiple nodes based on their descriptions.

        Args:
            description (str): The text description from which to infer node distributions.
            node_names (list[str]): A list of node names for which to suggest distributions.
            gpt_model (str, optional): The GPT model to use for suggestions. Defaults to 'gpt-4o-mini'.
            temperature (float, optional): The randomness of the model's output. Defaults to 0.

        Returns:
            dict[str, str]: A dictionary mapping each node name to its suggested distribution type.
        """
        node_distr_dict = {}
        for node_name in node_names:
            node_distr = self._suggest_vertex_distribution(
                description, node_name, gpt_model, temperature
            )
            node_distr_dict[node_name] = node_distr

        return node_distr_dict

__init__(openai_client)

Initializes the NodeDistributer with an OpenAI client.

Parameters:

Name Type Description Default
openai_client OpenAI | AsyncOpenAI

An instance of OpenAI or AsyncOpenAI client to interact with the API.

 required
Source code in code/NLI_suggest_node_distribution.py 
17
18
19
20
21
22
23
24
25
def __init__(self, openai_client: OpenAI | AsyncOpenAI) -> None:
    """
    Initializes the NodeDistributer with an OpenAI client.

    Args:
        openai_client (OpenAI | AsyncOpenAI): An instance of OpenAI or AsyncOpenAI client to interact with the API.
    """

    super(NodeDistributer, self).__init__(openai_client)

suggest_vertex_distributions(description, node_names, gpt_model='gpt-4o-mini', temperature=0)

Suggests distribution types for multiple nodes based on their descriptions.

Parameters:

Name Type Description Default
description str

The text description from which to infer node distributions.

 required
node_names list[str]

A list of node names for which to suggest distributions.

 required
gpt_model str

The GPT model to use for suggestions. Defaults to 'gpt-4o-mini'.

 'gpt-4o-mini'
temperature float

The randomness of the model's output. Defaults to 0.

 0

Returns:

Type Description
str

dict[str, str]: A dictionary mapping each node name to its suggested distribution type.
Source code in code/NLI_suggest_node_distribution.py 
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
def suggest_vertex_distributions(
    self,
    description: str,
    node_names: list[str],
    gpt_model: str = "gpt-4o-mini",
    temperature: float = 0,
) -> str:
    """
    Suggests distribution types for multiple nodes based on their descriptions.

    Args:
        description (str): The text description from which to infer node distributions.
        node_names (list[str]): A list of node names for which to suggest distributions.
        gpt_model (str, optional): The GPT model to use for suggestions. Defaults to 'gpt-4o-mini'.
        temperature (float, optional): The randomness of the model's output. Defaults to 0.

    Returns:
        dict[str, str]: A dictionary mapping each node name to its suggested distribution type.
    """
    node_distr_dict = {}
    for node_name in node_names:
        node_distr = self._suggest_vertex_distribution(
            description, node_name, gpt_model, temperature
        )
        node_distr_dict[node_name] = node_distr

    return node_distr_dict

Functions

Graph utils

Asynchronous Operations

The library supports asynchronous operations to improve efficiency when processing multiple requests. Ensure your environment supports asynchronous programming if you choose to use these features.