self-constructing_graph/agent/libs/graph.py

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