mseep-txtai 9.1.1__py3-none-any.whl

txtai/graph/factory.py

@@ -0,0 +1,61 @@
+"""
+Factory module
+"""
+
+from ..util import Resolver
+
+from .networkx import NetworkX
+from .rdbms import RDBMS
+
+
+class GraphFactory:
+    """
+    Methods to create graphs.
+    """
+
+    @staticmethod
+    def create(config):
+        """
+        Create a Graph.
+
+        Args:
+            config: graph configuration
+
+        Returns:
+            Graph
+        """
+
+        # Graph instance
+        graph = None
+        backend = config.get("backend", "networkx")
+
+        # Create graph instance
+        if backend == "networkx":
+            graph = NetworkX(config)
+        elif backend == "rdbms":
+            graph = RDBMS(config)
+        else:
+            graph = GraphFactory.resolve(backend, config)
+
+        # Store config back
+        config["backend"] = backend
+
+        return graph
+
+    @staticmethod
+    def resolve(backend, config):
+        """
+        Attempt to resolve a custom backend.
+
+        Args:
+            backend: backend class
+            config: index configuration parameters
+
+        Returns:
+            Graph
+        """
+
+        try:
+            return Resolver()(backend)(config)
+        except Exception as e:
+            raise ImportError(f"Unable to resolve graph backend: '{backend}'") from e

txtai/graph/networkx.py

@@ -0,0 +1,275 @@
+"""
+NetworkX module
+"""
+
+import os
+
+from tempfile import TemporaryDirectory
+
+# Conditional import
+try:
+    import networkx as nx
+
+    from networkx.algorithms.community import asyn_lpa_communities, greedy_modularity_communities, louvain_partitions
+    from networkx.readwrite import json_graph
+
+    NETWORKX = True
+except ImportError:
+    NETWORKX = False
+
+from ..archive import ArchiveFactory
+from ..serialize import SerializeError, SerializeFactory
+
+from .base import Graph
+from .query import Query
+
+
+# pylint: disable=R0904
+class NetworkX(Graph):
+    """
+    Graph instance backed by NetworkX.
+    """
+
+    def __init__(self, config):
+        super().__init__(config)
+
+        if not NETWORKX:
+            raise ImportError('NetworkX is not available - install "graph" extra to enable')
+
+    def create(self):
+        return nx.Graph()
+
+    def count(self):
+        return self.backend.number_of_nodes()
+
+    def scan(self, attribute=None, data=False):
+        # Full graph
+        graph = self.backend
+
+        # Filter graph to nodes having a specified attribute
+        if attribute:
+            graph = nx.subgraph_view(self.backend, filter_node=lambda x: attribute in self.node(x))
+
+        # Return either list of matching ids or tuple of (id, attribute dictionary)
+        return graph.nodes(data=True) if data else graph
+
+    def node(self, node):
+        return self.backend.nodes.get(node)
+
+    def addnode(self, node, **attrs):
+        self.backend.add_node(node, **attrs)
+
+    def addnodes(self, nodes):
+        self.backend.add_nodes_from(nodes)
+
+    def removenode(self, node):
+        if self.hasnode(node):
+            self.backend.remove_node(node)
+
+    def hasnode(self, node):
+        return self.backend.has_node(node)
+
+    def attribute(self, node, field):
+        return self.node(node).get(field) if self.hasnode(node) else None
+
+    def addattribute(self, node, field, value):
+        if self.hasnode(node):
+            self.node(node)[field] = value
+
+    def removeattribute(self, node, field):
+        return self.node(node).pop(field, None) if self.hasnode(node) else None
+
+    def edgecount(self):
+        return self.backend.number_of_edges()
+
+    def edges(self, node):
+        edges = self.backend.adj.get(node)
+        if edges:
+            return dict(sorted(edges.items(), key=lambda x: x[1].get("weight", 0), reverse=True))
+
+        return None
+
+    def addedge(self, source, target, **attrs):
+        self.backend.add_edge(source, target, **attrs)
+
+    def addedges(self, edges):
+        self.backend.add_edges_from(edges)
+
+    def hasedge(self, source, target=None):
+        if target is None:
+            edges = self.backend.adj.get(source)
+            return len(edges) > 0 if edges else False
+
+        return self.backend.has_edge(source, target)
+
+    def centrality(self):
+        rank = nx.degree_centrality(self.backend)
+        return dict(sorted(rank.items(), key=lambda x: x[1], reverse=True))
+
+    def pagerank(self):
+        rank = nx.pagerank(self.backend, weight="weight")
+        return dict(sorted(rank.items(), key=lambda x: x[1], reverse=True))
+
+    def showpath(self, source, target):
+        # pylint: disable=E1121
+        return nx.shortest_path(self.backend, source, target, self.distance)
+
+    def isquery(self, queries):
+        return Query().isquery(queries)
+
+    def parse(self, query):
+        return Query().parse(query)
+
+    def search(self, query, limit=None, graph=False):
+        # Run graph query
+        results = Query()(self, query, limit)
+
+        # Transform into filtered graph
+        if graph:
+            nodes = set()
+            for column in results.values():
+                for value in column:
+                    if isinstance(value, list):
+                        # Path group
+                        nodes.update([node for node in value if node and not isinstance(node, dict)])
+                    elif isinstance(value, dict):
+                        # Nodes by id attribute
+                        nodes.update(uid for uid, attr in self.scan(data=True) if attr["id"] == value["id"])
+                    elif value is not None:
+                        # Single node id
+                        nodes.add(value)
+
+            return self.filter(list(nodes))
+
+        # Transform columnar structure into rows
+        keys = list(results.keys())
+        rows, count = [], len(results[keys[0]])
+
+        for x in range(count):
+            rows.append({str(key): results[key][x] for key in keys})
+
+        return rows
+
+    def communities(self, config):
+        # Get community detection algorithm
+        algorithm = config.get("algorithm")
+
+        if algorithm == "greedy":
+            communities = greedy_modularity_communities(self.backend, weight="weight", resolution=config.get("resolution", 100))
+        elif algorithm == "lpa":
+            communities = asyn_lpa_communities(self.backend, weight="weight", seed=0)
+        else:
+            communities = self.louvain(config)
+
+        return communities
+
+    def load(self, path):
+        try:
+            # Load graph data
+            data = SerializeFactory.create().load(path)
+
+            # Add data to graph
+            self.backend = self.create()
+            self.backend.add_nodes_from(data["nodes"])
+            self.backend.add_edges_from(data["edges"])
+
+            # Load categories
+            self.categories = data.get("categories")
+
+            # Load topics
+            self.topics = data.get("topics")
+
+        except SerializeError:
+            # Backwards compatible support for legacy TAR format
+            self.loadtar(path)
+
+    def save(self, path):
+        # Save graph data
+        SerializeFactory.create().save(
+            {
+                "nodes": [(uid, self.node(uid)) for uid in self.scan()],
+                "edges": list(self.backend.edges(data=True)),
+                "categories": self.categories,
+                "topics": self.topics,
+            },
+            path,
+        )
+
+    def loaddict(self, data):
+        self.backend = json_graph.node_link_graph(data, name="indexid")
+        self.categories, self.topics = data.get("categories"), data.get("topics")
+
+    def savedict(self):
+        data = json_graph.node_link_data(self.backend, name="indexid")
+        data["categories"] = self.categories
+        data["topics"] = self.topics
+
+        return data
+
+    def louvain(self, config):
+        """
+        Runs the Louvain community detection algorithm.
+
+        Args:
+            config: topic configuration
+
+        Returns:
+            list of [ids] per community
+        """
+
+        # Partition level to use
+        level = config.get("level", "best")
+
+        # Run community detection
+        results = list(louvain_partitions(self.backend, weight="weight", resolution=config.get("resolution", 100), seed=0))
+
+        # Get partition level (first or best)
+        return results[0] if level == "first" else results[-1]
+
+    # pylint: disable=W0613
+    def distance(self, source, target, attrs):
+        """
+        Computes distance between source and target nodes using weight.
+
+        Args:
+            source: source node
+            target: target node
+            attrs: edge attributes
+
+        Returns:
+            distance between source and target
+        """
+
+        # Distance is 1 - score. Skip minimal distances as they are near duplicates.
+        distance = max(1.0 - attrs["weight"], 0.0)
+        return distance if distance >= 0.15 else 1.00
+
+    def loadtar(self, path):
+        """
+        Loads a graph from the legacy TAR file.
+
+        Args:
+            path: path to graph
+        """
+
+        # Pickle serialization - backwards compatible
+        serializer = SerializeFactory.create("pickle")
+
+        # Extract files to temporary directory and load content
+        with TemporaryDirectory() as directory:
+            # Unpack files
+            archive = ArchiveFactory.create(directory)
+            archive.load(path, "tar")
+
+            # Load graph backend
+            self.backend = serializer.load(f"{directory}/graph")
+
+            # Load categories, if necessary
+            path = f"{directory}/categories"
+            if os.path.exists(path):
+                self.categories = serializer.load(path)
+
+            # Load topics, if necessary
+            path = f"{directory}/topics"
+            if os.path.exists(path):
+                self.topics = serializer.load(path)

txtai/graph/query.py

@@ -0,0 +1,181 @@
+"""
+Query module
+"""
+
+import logging
+import re
+
+try:
+    from grandcypher import GrandCypher
+
+    GRANDCYPHER = True
+except ImportError:
+    GRANDCYPHER = False
+
+# Logging configuration
+logger = logging.getLogger(__name__)
+
+
+class Query:
+    """
+    Runs openCypher graph queries using the GrandCypher library. This class also supports search functions.
+    """
+
+    # Similar token
+    SIMILAR = "__SIMILAR__"
+
+    def __init__(self):
+        """
+        Create a new graph query instance.
+        """
+
+        if not GRANDCYPHER:
+            raise ImportError('GrandCypher is not available - install "graph" extra to enable')
+
+    def __call__(self, graph, query, limit):
+        """
+        Runs a graph query.
+
+        Args:
+            graph: graph instance
+            query: graph query, can be a full query string or a parsed query dictionary
+            limit: number of results
+
+        Returns:
+            results
+        """
+
+        # Results by attribute and ids filter
+        attributes, uids = None, None
+
+        # Build the query from a parsed query
+        if isinstance(query, dict):
+            query, attributes, uids = self.build(query)
+
+        # Filter graph, if applicable
+        if uids:
+            graph = self.filter(graph, attributes, uids)
+
+        # Debug log graph query
+        logger.debug(query)
+
+        # Run openCypher query
+        return GrandCypher(graph.backend, limit if limit else 3).run(query)
+
+    def isquery(self, queries):
+        """
+        Checks a list of queries to see if all queries are openCypher queries.
+
+        Args:
+            queries: list of queries to check
+
+        Returns:
+            True if all queries are openCypher queries
+        """
+
+        # Check for required graph query clauses
+        return all(query and query.strip().startswith("MATCH ") and "RETURN " in query for query in queries)
+
+    def parse(self, query):
+        """
+        Parses a graph query. This method supports parsing search functions and replacing them with placeholders.
+
+        Args:
+            query: graph query
+
+        Returns:
+            parsed query as a dictionary
+        """
+
+        # Parameters
+        where, limit, nodes, similar = None, None, [], []
+
+        # Parse where clause
+        match = re.search(r"where(.+?)return", query, flags=re.DOTALL | re.IGNORECASE)
+        if match:
+            where = match.group(1).strip()
+
+        # Parse limit clause
+        match = re.search(r"limit\s+(\d+)", query, flags=re.DOTALL | re.IGNORECASE)
+        if match:
+            limit = match.group(1)
+
+        # Parse similar clauses
+        for x, match in enumerate(re.finditer(r"similar\((.+?)\)", query, flags=re.DOTALL | re.IGNORECASE)):
+            # Replace similar clause with placeholder
+            query = query.replace(match.group(0), f"{Query.SIMILAR}{x}")
+
+            # Parse similar clause parameters
+            params = [param.strip().replace("'", "").replace('"', "") for param in match.group(1).split(",")]
+            nodes.append(params[0])
+            similar.append(params[1:])
+
+        # Return parsed query
+        return {
+            "query": query,
+            "where": where,
+            "limit": limit,
+            "nodes": nodes,
+            "similar": similar,
+        }
+
+    def build(self, parse):
+        """
+        Constructs a full query from a parsed query. This method supports substituting placeholders with search results.
+
+        Args:
+            parse: parsed query
+
+        Returns:
+            graph query
+        """
+
+        # Get query. Initialize attributes and uids.
+        query, attributes, uids = parse["query"], {}, {}
+
+        # Replace similar clause with id query
+        if "results" in parse:
+            for x, result in enumerate(parse["results"]):
+                # Get query node
+                node = parse["nodes"][x]
+
+                # Add similar match attribute
+                attribute = f"match_{x}"
+                clause = f"{node}.{attribute} > 0"
+
+                # Replace placeholder with earch results
+                query = query.replace(f"{Query.SIMILAR}{x}", f"{clause}")
+
+                # Add uids and scores
+                for uid, score in result:
+                    if uid not in uids:
+                        uids[uid] = score
+
+                # Add results by attribute matched
+                attributes[attribute] = result
+
+        # Return query, results by attribute matched and ids filter
+        return query, attributes, uids.items()
+
+    def filter(self, graph, attributes, uids):
+        """
+        Filters the input graph by uids. This method also adds similar match attributes.
+
+        Args:
+            graph: graph instance
+            attributes: results by attribute matched
+            uids: single list with all matching ids
+
+        Returns:
+            filtered graph
+        """
+
+        # Filter the graph
+        graph = graph.filter(uids)
+
+        # Add similar match attributes
+        for attribute, result in attributes.items():
+            for uid, score in result:
+                graph.addattribute(uid, attribute, score)
+
+        return graph

txtai/graph/rdbms.py

@@ -0,0 +1,113 @@
+"""
+RDBMS module
+"""
+
+import os
+
+# Conditional import
+try:
+    from grand import Graph
+    from grand.backends import SQLBackend, InMemoryCachedBackend
+
+    from sqlalchemy import create_engine, text, StaticPool
+    from sqlalchemy.schema import CreateSchema
+
+    ORM = True
+except ImportError:
+    ORM = False
+
+from .networkx import NetworkX
+
+
+class RDBMS(NetworkX):
+    """
+    Graph instance backed by a relational database.
+    """
+
+    def __init__(self, config):
+        # Check before super() in case those required libraries are also not available
+        if not ORM:
+            raise ImportError('RDBMS is not available - install "graph" extra to enable')
+
+        super().__init__(config)
+
+        # Graph and database instances
+        self.graph = None
+        self.database = None
+
+    def __del__(self):
+        if hasattr(self, "database") and self.database:
+            self.database.close()
+
+    def create(self):
+        # Create graph instance
+        self.graph, self.database = self.connect()
+
+        # Clear previous graph, if available
+        for table in [self.config.get("nodes", "nodes"), self.config.get("edges", "edges")]:
+            self.database.execute(text(f"DELETE FROM {table}"))
+
+        # Return NetworkX compatible backend
+        return self.graph.nx
+
+    def scan(self, attribute=None, data=False):
+        if attribute:
+            for node in self.backend:
+                attributes = self.node(node)
+                if attribute in attributes:
+                    yield (node, attributes) if data else node
+        else:
+            yield from super().scan(attribute, data)
+
+    def load(self, path):
+        # Create graph instance
+        self.graph, self.database = self.connect()
+
+        # Store NetworkX compatible backend
+        self.backend = self.graph.nx
+
+    def save(self, path):
+        self.database.commit()
+
+    def close(self):
+        # Parent logic
+        super().close()
+
+        # Close database connection
+        self.database.close()
+
+    def filter(self, nodes, graph=None):
+        return super().filter(nodes, graph if graph else NetworkX(self.config))
+
+    def connect(self):
+        """
+        Connects to a graph backed by a relational database.
+
+        Args:
+            Graph database instance
+        """
+
+        # Keyword arguments for SQLAlchemy
+        kwargs = {"poolclass": StaticPool, "echo": False}
+        url = self.config.get("url", os.environ.get("GRAPH_URL"))
+
+        # Set default schema, if necessary
+        schema = self.config.get("schema")
+        if schema:
+            # Check that schema exists
+            engine = create_engine(url)
+            with engine.begin() as connection:
+                connection.execute(CreateSchema(schema, if_not_exists=True) if "postgresql" in url else text("SELECT 1"))
+
+            # Set default schema
+            kwargs["connect_args"] = {"options": f'-c search_path="{schema}"'} if "postgresql" in url else {}
+
+        backend = SQLBackend(
+            db_url=url,
+            node_table_name=self.config.get("nodes", "nodes"),
+            edge_table_name=self.config.get("edges", "edges"),
+            sqlalchemy_kwargs=kwargs,
+        )
+
+        # pylint: disable=W0212
+        return Graph(backend=InMemoryCachedBackend(backend, maxsize=None)), backend._connection