self-constructing_graph/agent/router/kb_router.py

import sys,os
current_path = os.getcwd()
sys.path.append(current_path)

from config.site import SiteConfig
from fastapi import APIRouter, Depends, Query
from db.database import get_db
from sqlalchemy.orm import Session
from agent.models.web.response import StandardResponse,FAILED,SUCCESS
from agent.models.web.request import BasicRequest
from agent.libs.graph import GraphBusiness
from agent.libs.auth import verify_session_id, SessionValues
import logging

router = APIRouter(prefix="/kb", tags=["knowledge build interface"])
logger = logging.getLogger(__name__)
config = SiteConfig()

LOG_DIR = config.get_config("TASK_LOG_DIR", current_path)

    # job_category = Column(String(64), nullable=False)
    # job_name = Column(String(64))
    # job_details = Column(Text, nullable=False)
    # job_creator = Column(String(64), nullable=False)
    # job_logs = Column(Text, nullable=True)
    # job_files = Column(String(300), nullable=True)
@router.post('/summary', response_model=StandardResponse)
def summary_func(request:BasicRequest, db: Session = Depends(get_db), sess:SessionValues = Depends(verify_session_id))->StandardResponse:
    if request.action != "get_summary":
        return StandardResponse(code=FAILED, message="invalid action")
    graph_id = request.get_param("graph_id",0)
    biz = GraphBusiness(db)
    summary = biz.get_graph_summary(graph_id=graph_id)
    if summary:
        logger.info(summary)
        return StandardResponse(code=SUCCESS, message="summary found", records=[summary])
    else:
        return StandardResponse(code=FAILED, message="summary not found",records=[])

@router.post('/schemas', response_model=StandardResponse)
def schemas_func(request:BasicRequest, db: Session = Depends(get_db), sess:SessionValues = Depends(verify_session_id))->StandardResponse:
    if request.action== "get_nodes_schemas":
        graph_id = request.get_param("graph_id",0)
        biz = GraphBusiness(db)
        schemas = biz.get_nodes_categories(graph_id=graph_id)
        if schemas:
            return StandardResponse(code=SUCCESS, message="schemas found", records=schemas)
    if request.action== "get_edges_schemas":
        graph_id = request.get_param("graph_id",0)
        biz = GraphBusiness(db)
        schemas = biz.get_edges_categories(graph_id=graph_id)
        if schemas:
            return StandardResponse(code=SUCCESS, message="schemas found", records=schemas)
    return StandardResponse(code=FAILED, message="invalid action")

@router.post('/nodes', response_model=StandardResponse)
def nodes_func(request:BasicRequest, db: Session = Depends(get_db), sess:SessionValues = Depends(verify_session_id))->StandardResponse:    
    if (request.action == "search_nodes"):
        node_name = request.get_param("name","")
        category = request.get_param("category","")
        graph_id = request.get_param("graph_id",0)   
        biz = GraphBusiness(db)   
        if node_name == "":
            return StandardResponse(code=FAILED, message="node name is empty", records=[])
        if category == "":
            return StandardResponse(code=FAILED, message="category is empty", records=[])
        if graph_id == 0:
            return StandardResponse(code=FAILED, message="graph id is empty", records=[])
        
        nodes = biz.search_like_node_by_name(graph_id=graph_id, category=category, name=node_name)
        if nodes:
            return StandardResponse(code=SUCCESS, message="nodes found", records=nodes)
        
        else:
            return StandardResponse(code=FAILED, message="search job failed")
    elif (request.action == "get_nodes"):
        graph_id = request.get_param("graph_id",0)
        page = request.get_param("page",1)
        page_size = request.get_param("page_size",1)
        biz = GraphBusiness(db)
        nodes = biz.get_nodes_by_page(graph_id=graph_id, page=page, page_size=page_size)
        if nodes:
            return StandardResponse(code=SUCCESS, message="nodes found", records=nodes)
    elif (request.action == "neighbors"):
        graph_id = request.get_param("graph_id",0)
        node_id = request.get_param("node_id",0)
        if node_id>0:
            biz = GraphBusiness(db)
            node = biz.get_node_by_id(graph_id=graph_id, node_id=node_id)
            if node is None:
                return StandardResponse(code=FAILED, message="node not found", records=[])
            nodes_in = biz.get_neighbors(graph_id=graph_id, node_id=node_id, direction="in")
            nodes_out = biz.get_neighbors(graph_id=graph_id, node_id=node_id, direction="out")
            nodes_all = []
            nodes_all.append({"id": node.id, "name":node.name, "category":node.category, "direction":"self"})
            for node in nodes_in:
                nodes_all.append({"id": node.id, "name":node.name, "category":node.category, "direction":"in"})
            for node in nodes_out:
                nodes_all.append({"id": node.id, "name":node.name, "category":node.category, "direction":"out"})
            
            return StandardResponse(code=SUCCESS, message="nodes found", records=nodes_all)
    return StandardResponse(code=FAILED, message="invalid action")
        
        
        
kb_router = router