kg_person/KG_Search_API.py

from pprint import pprint
import requests
import json
from hashlib import md5
import uvicorn
from fastapi import FastAPI
from typing import Optional

app = FastAPI()

def getQaAttachment(data):
    nodes = []
    ids = []
    links = []
    for d in data:
        if not d['target']['id']:
            d['target']['id'] = md5(str(d['target']['entity']).encode(encoding='UTF-8')).hexdigest()
        if not d['target']['id'] in ids:
            nodes.append({"id":d['target']['id'], "name": str(d['target']['entity']), "target":True, "source":False})
            ids.append(d['target']['id'])
        if not d['source']['id'] in ids:
            nodes.append({"id":d['source']['id'], "name": d['source']['entity'], "source":True, "target":False})
            ids.append(d['source']['id'])
        links.append({'source': d['source']['id'], 'name':d["relation"], 'target': d['target']['id']})
    return nodes, links

# 请求头
headers = {
	"token": '369ca613c7a74a58b0f95be2cfd59257',
	"Content-Type": "application/json",
}



@app.get('/person_to_person')
def search_person(query: str, less: Optional[bool] = False):
	"""
	# 人与人的相关关系查询
	# 必选参数:
	# query: 人名
	# 可选参数
	# less: 显示少量关系
	"""
	url = 'http://180.76.188.39:8085/mpks/api/extra/gremlin'
	if less:
		json_obj = {
			# "gremlin": "g.has('name.@value', MATCH, '"+ query +"').outE(['当前公司','最高学历学校']).inV.inE('相关机构').outV.with('*').graph"
			"gremlin": "g.has('name.@value', MATCH, '"+ query +"').outE('相关机构').limit(8).inV.inE('相关机构').outV.with('*').graph"
		}
	else:
		json_obj = {
			"gremlin": "g.has('name.@value', MATCH, '"+ query +"').outE('相关机构').inV.inE('相关机构').outV.with('*').graph"
		}
	# 请求查询
	r = requests.post(url, headers=headers, json=json_obj)
	rst = json.loads(r.text)

	entity_array = []
	nodes = []
	links = []
	ids = dict([])
	
	# 查询成功
	if (rst['errno'] == 0) and rst['data']:
		# 所有实体
		for entity in rst['data']['vertices']:
			if entity["@id"] not in ids:
				ids[entity["@id"]] = entity["name"]
				nodes.append({
					"id": entity["@id"],
					"name": entity["name"],
					"target": False if (entity["@type"] == "人才特征demo") else True,
					"source": True if (entity["@type"] == "人才特征demo") else False,
					"itemStyle": {
            			"normal": {
              				"color": 'red' if (entity["@type"] == "人才特征demo") else 'blue'
            			}
            		},
            	})
		for edge in rst['data']['edges']:
			source = {'id': edge['@from'], 'entity': ids[edge['@from']], "refId": None, "entityIndex": None}
			target = {'id': edge['@to'], 'entity': ids[edge['@to']], "refId": None, "entityIndex": None}
			relation = edge['@label']
			entity_array.append({"source": source, "relation": relation, "target": target})
			links.append({'source': edge['@from'], 'name': edge['@label'], 'target': edge['@to']})
	else:
		pprint(rst['errno'])
		pprint(rst['msg'])
	return {"errno": 0, "msg": [], "graph": entity_array, "nodes": nodes, "links": links}



@app.get('/search_query')
def search_query(query: str):
	"""
	# 搜索接口,模糊搜索
	# 參數：
	# query: 查詢的關鍵詞
	"""
	entity_array = []
	base_url = 'http://180.76.188.39:8085/mpks/api/search'
	json_obj = {
		"query": query,
		"sort": "relevance",
		"needCorrect": True,
		"saveHistory": False
	}
	r = requests.post(base_url, headers=headers, json=json_obj)
	rst = json.loads(r.text)
	# 查询结果
	if (rst['errno'] == 0) and rst['data']:
		# 当前实体
		source = {'id': rst['data']['results'][0]['entityList'][0]['@id'], 'entity': rst['data']['results'][0]['entityList'][0]['name'],'refId': None, "entityIndex": None}
		# 实体属性关系
		for dic in rst['data']['results'][0]['entityList'][0]['properties']:
			relation = dic['key']
			# 关系
			if isinstance(dic['value'], dict) and ("@id" in dic['value']):
				target = {"id": dic['value']['@id'], "entity":dic['value']['value'], 'refId': None, "entityIndex": None}
				entity_array.append({"source": source, "relation": relation, "target": target})
			# 属性
			elif (not isinstance(dic['value'], list)) or ("@id" not in dic['value'][0]):
				target = {"id": None, "entity":dic['value'], 'refId': None, "entityIndex": None}
				entity_array.append({"source": source, "relation": relation, "target": target})
			# 关系
			else:
				for vec in dic['value']:
					target = {"id": vec["@id"], "entity": vec["value"], 'refId': None, "entityIndex": None}
					entity_array.append({"source": source, "relation": relation, "target": target})
		# 边
		for edge in rst['data']['results'][0]['entityList'][0]['graphData']['vertices']:
			relation = edge['@type'].replace('demo', '')
			target = {"id": edge["@id"], "entity": edge["name"], 'refId': None, "entityIndex": None}
			entity_array.append({"source": source, "relation": relation, "target": target})
	else:
		pprint(rst['errno'])
		pprint(rst['msg'])
	nodes, links = getQaAttachment(entity_array)
	return {"errno": 0, "msg": [], "graph": entity_array, "nodes": nodes, "links": links}


@app.get('/search_gremlin')
def search_gremlin(query: Optional[str] = None, _id: Optional[str] = None):
	"""
	# 搜索接口,精确搜索
	# 參數：
	# query[可选]: 使用实体名称查询
	# _id  [可选]: 使用实体 id 查询
	# 同时填写时 _id 优先
	g.has("name.@value",MATCH,"于策").outE("相关机构").inV.inE("相关机构").outV.with("*").graph
	"""
	# 请求地址
	url = 'http://180.76.188.39:8085/mpks/api/extra/gremlin'
	# 查询语句
	if _id:
		json_obj = {
			"gremlin" : "g.key('"+ _id +"').both.with('*').graph"
		}
	elif query:
		json_obj = {
			"gremlin" : "g.has('name.@value', MATCH, '"+ query +"').both.with('*').graph"
		}
	else:
		return {"errno": 3001, "msg": "can not get query or id", "data": []}
	# 请求查询
	r = requests.post(url, headers=headers, json=json_obj)
	rst = json.loads(r.text)

	entity_array = []
	
	# 查询成功
	if (rst['errno'] == 0) and rst['data']:
		# 当前实体
		source = {'id': rst['data']['vertices'][0]['@id'], 'entity': rst['data']['vertices'][0]['name'],'refId': None, "entityIndex": None}
		# 属性
		for relation in rst['data']['vertices'][0].keys():
			if relation not in ['@context', '@del', '@edge_number', '@formattype', '@fromtype', '@fromurl', '@id', '@kbid', '@nodeid', '@semiid', '@tags', '@type', '_id', '_type', 'alias', 'appId', 'name', 'nodeId', 'tags', '教育经历', '工作经历', '项目经历', '培训和海外经历', '最高学历学校', '当前公司']:
				target = {'id': None, 'entity': rst['data']['vertices'][0][relation], 'refId': None, "entityIndex": None}
				if rst['data']['vertices'][0][relation]:
					entity_array.append({"source": source, "relation": relation, "target": target})
		# 关系
		if len(rst['data']['vertices']) > 1:
			for index in range(1, len(rst['data']['vertices'])):
				target = {'id': rst['data']['vertices'][index]['@id'], 'entity': rst['data']['vertices'][index]['name'], 'refId': None, "entityIndex": None}
				for edge in rst['data']['edges']:
					if edge['@from'] == rst['data']['vertices'][index]['@id']:
						entity_array.append({"source": target, "relation": edge['@label'].replace('demo', ''), "target": source})
						break
				else:
					entity_array.append({"source": source, "relation": rst['data']['vertices'][index]['@type'].replace('demo', ''), "target": target})
		# 边
		# for edge in rst['data']['edges']:
		# 	target = {'id': edge['@to'], 'entity': None, "refId": edge['@from'], "entityIndex": source_index}
		# 	relation = edge['@label']
		# 	entity_array.append({"source": source, "relation": relation, "target": target})
	else:
		pprint(rst['errno'])
		pprint(rst['msg'])
	nodes, links = getQaAttachment(entity_array)
	return {"errno": 0, "msg": [], "graph": entity_array, "nodes": nodes, "links": links}


if __name__ == '__main__':
	uvicorn.run(app=app, host='0.0.0.0', port=9000)