import uuid import os import subprocess import logging logger = logging.getLogger(__name__) from datetime import datetime from agent.models.db.graph import DbKgGraphs,DbKgNode,DbKgEdge,DbKgProp,DbKgEdgeProp from sqlalchemy.orm import load_only class GraphBusiness: def __init__(self, db): self.db = db self.graphs = {} def create_graph(self, category:str, name: str, description: str, graph_settings:str): graph = DbKgGraphs(category=category,name=name, description=description, graph_settings=graph_settings) self.db.add(graph) self.db.commit() self.db.refresh(graph) return graph def get_graph(self, graph_id: int): graph = self.db.query(DbKgGraphs).filter(DbKgGraphs.id == graph_id).first() return graph def create_node(self, graph_id: int, name: str, category: str, props: dict): node = self.db.query(DbKgNode).filter(DbKgNode.graph_id == graph_id, DbKgNode.name == name, DbKgNode.category == category).first() if node: return node node = DbKgNode(graph_id=graph_id, name=name, category=category) self.db.add(node) self.db.commit() self.db.refresh(node) for key, value in props.items(): self.create_node_prop(category=1, ref_id=node.id, prop_name=key, prop_value=value, prop_title=key) self.db.commit() return node def create_edge(self, graph_id: int, src_id: int, dest_id: int, category: str, name: str, props: dict): edge = self.db.query(DbKgEdge).filter(DbKgEdge.graph_id == graph_id, DbKgEdge.src_id == src_id, DbKgEdge.dest_id == dest_id, DbKgEdge.category == category, DbKgEdge.name == name).first() if edge: for key, value in props.items(): prop = self.db.query(DbKgEdgeProp).filter(DbKgEdgeProp.ref_id == edge.id, DbKgEdgeProp.prop_name == key).first() if prop: continue self.create_edge_prop(category=1, ref_id=edge.id, prop_name=key, prop_value=value, prop_title=key) return edge edge = DbKgEdge(graph_id=graph_id, src_id=src_id, dest_id=dest_id, category=category, name=name) self.db.add(edge) self.db.commit() self.db.refresh(edge) for key, value in props.items(): self.create_edge_prop(category=1, ref_id=edge.id, prop_name=key, prop_value=value, prop_title=key) self.db.commit() return edge def create_node_prop(self, category:str, ref_id: int, prop_name: str, prop_value: str, prop_title:str, commit=False): prop = DbKgProp(category=category, ref_id=ref_id, prop_name=prop_name, prop_value=prop_value, prop_title=prop_title) self.db.add(prop) if commit: self.db.commit() self.db.refresh(prop) return prop def create_edge_prop(self, category:str, ref_id: int, prop_name: str, prop_value: str, prop_title:str, commit=False): prop = DbKgEdgeProp(category=category, ref_id=ref_id, prop_name=prop_name, prop_value=prop_value, prop_title=prop_title) self.db.add(prop) if commit: self.db.commit() self.db.refresh(prop) return prop def search_like_node_by_name(self, name: str, category: str, graph_id: int): if name == "": name = "%" if category == "" or category == "any": nodes = self.db.query(DbKgNode).filter(DbKgNode.name.like(f"%{name}%"), DbKgNode.graph_id == graph_id).limit(100).all() return nodes nodes = self.db.query(DbKgNode)\ .filter(DbKgNode.name.like(f"%{name}%"), DbKgNode.category == category, DbKgNode.graph_id == graph_id).limit(100).all() return nodes def get_nodes_by_page(self, graph_id: int, page: int, page_size: int): nodes = self.db.query(DbKgNode).filter(DbKgNode.graph_id == graph_id).limit(page_size).offset((page - 1) * page_size).all() return nodes def get_nodes_count(self, graph_id: int): count = self.db.query(DbKgNode).filter(DbKgNode.graph_id == graph_id).count() return count def get_neighbors(self, graph_id: int, node_id: int, direction: str): nodes = [] if direction == "in": edges = self.db.query(DbKgEdge).filter(DbKgEdge.graph_id == graph_id, DbKgEdge.dest_id == node_id).limit(100).all() for edge in edges: nodes.append(edge.src_node) elif direction == "out": edges = self.db.query(DbKgEdge).filter(DbKgEdge.graph_id == graph_id, DbKgEdge.src_id == node_id).limit(100).all() for edge in edges: nodes.append(edge.dest_node) return nodes def get_node_by_id(self, graph_id: int, node_id: int): node = self.db.query(DbKgNode).filter(DbKgNode.graph_id == graph_id, DbKgNode.id == node_id).first() return node def get_graph_summary(self, graph_id: int): nodes_count = self.db.query(DbKgNode).filter(DbKgNode.graph_id == graph_id).count() edges_count = self.db.query(DbKgEdge).filter(DbKgEdge.graph_id == graph_id).count() nodes_categorys = self.db.query(DbKgNode.category).filter(DbKgNode.graph_id == graph_id).distinct().all() edges_categorys = self.db.query(DbKgEdge.category).filter(DbKgEdge.graph_id == graph_id).distinct().all() nodes_categorys = [category[0] for category in nodes_categorys] edges_categorys = [category[0] for category in edges_categorys] return {"nodes_count": nodes_count, "edges_count": edges_count, "nodes_categories": nodes_categorys, "edges_categories": edges_categorys} def get_nodes_categories(self, graph_id: int): nodes_categorys = self.db.query(DbKgNode.category).filter(DbKgNode.graph_id == graph_id).distinct().all() nodes_categorys = [{'name':category[0]} for category in nodes_categorys] return nodes_categorys def get_edges_categories(self, graph_id: int): edges_categorys = self.db.query(DbKgEdge.category).filter(DbKgEdge.graph_id == graph_id).distinct().all() edges_categorys = [{'name':category[0]} for category in edges_categorys] return edges_categorys