|
|
@@ -14,14 +14,29 @@ from knowhub.knowhub_db.cascade import cascade_delete
|
|
|
|
|
|
load_dotenv()
|
|
|
|
|
|
-# 关联字段的子查询(从 junction table 读取,返回 JSON 数组)
|
|
|
+# 关联字段的子查询(从 junction table 读取)
|
|
|
+# 对于带 relation_type 的 *_knowledge 边,同时暴露两种视图:
|
|
|
+# - *_ids : 扁平 ID 列表(向后兼容,不含 type)
|
|
|
+# - *_links : [{id, relation_type}](含 type)
|
|
|
_REL_SUBQUERIES = """
|
|
|
(SELECT COALESCE(json_agg(rk.requirement_id), '[]'::json)
|
|
|
FROM requirement_knowledge rk WHERE rk.knowledge_id = knowledge.id) AS requirement_ids,
|
|
|
+ (SELECT COALESCE(json_agg(json_build_object(
|
|
|
+ 'id', rk2.requirement_id, 'relation_type', rk2.relation_type
|
|
|
+ )), '[]'::json)
|
|
|
+ FROM requirement_knowledge rk2 WHERE rk2.knowledge_id = knowledge.id) AS requirement_links,
|
|
|
(SELECT COALESCE(json_agg(ck.capability_id), '[]'::json)
|
|
|
FROM capability_knowledge ck WHERE ck.knowledge_id = knowledge.id) AS capability_ids,
|
|
|
+ (SELECT COALESCE(json_agg(json_build_object(
|
|
|
+ 'id', ck2.capability_id, 'relation_type', ck2.relation_type
|
|
|
+ )), '[]'::json)
|
|
|
+ FROM capability_knowledge ck2 WHERE ck2.knowledge_id = knowledge.id) AS capability_links,
|
|
|
(SELECT COALESCE(json_agg(tk.tool_id), '[]'::json)
|
|
|
FROM tool_knowledge tk WHERE tk.knowledge_id = knowledge.id) AS tool_ids,
|
|
|
+ (SELECT COALESCE(json_agg(json_build_object(
|
|
|
+ 'id', tk2.tool_id, 'relation_type', tk2.relation_type
|
|
|
+ )), '[]'::json)
|
|
|
+ FROM tool_knowledge tk2 WHERE tk2.knowledge_id = knowledge.id) AS tool_links,
|
|
|
(SELECT COALESCE(json_agg(kr.resource_id), '[]'::json)
|
|
|
FROM knowledge_resource kr WHERE kr.knowledge_id = knowledge.id) AS resource_ids,
|
|
|
(SELECT COALESCE(json_agg(json_build_object(
|
|
|
@@ -44,6 +59,26 @@ _SELECT_FIELDS = f"{_BASE_FIELDS}, {_REL_SUBQUERIES}"
|
|
|
_SELECT_FIELDS_WITH_EMB = f"task_embedding, content_embedding, {_SELECT_FIELDS}"
|
|
|
|
|
|
|
|
|
+def _normalize_links(data: Dict, links_key: str, ids_key: str, default_type: str):
|
|
|
+ """
|
|
|
+ 统一两种输入格式:
|
|
|
+ - {links_key: [{id, relation_type}, ...]} → 使用指定 type
|
|
|
+ - {ids_key: [id1, id2, ...]} → 使用 default_type
|
|
|
+ 两个 key 都没有返回 None(不更新)
|
|
|
+ """
|
|
|
+ if links_key in data and data[links_key] is not None:
|
|
|
+ out = []
|
|
|
+ for item in data[links_key]:
|
|
|
+ if isinstance(item, dict):
|
|
|
+ out.append((item['id'], item.get('relation_type', default_type)))
|
|
|
+ else:
|
|
|
+ out.append((item, default_type))
|
|
|
+ return out
|
|
|
+ if ids_key in data and data[ids_key] is not None:
|
|
|
+ return [(i, default_type) for i in data[ids_key]]
|
|
|
+ return None
|
|
|
+
|
|
|
+
|
|
|
class PostgreSQLStore:
|
|
|
def __init__(self):
|
|
|
"""初始化 PostgreSQL 连接"""
|
|
|
@@ -54,7 +89,7 @@ class PostgreSQLStore:
|
|
|
password=os.getenv('KNOWHUB_PASSWORD'),
|
|
|
database=os.getenv('KNOWHUB_DB_NAME')
|
|
|
)
|
|
|
- self.conn.autocommit = False
|
|
|
+ self.conn.autocommit = True
|
|
|
print(f"[PostgreSQL] 已连接到远程数据库: {os.getenv('KNOWHUB_DB')}")
|
|
|
|
|
|
def _reconnect(self):
|
|
|
@@ -65,7 +100,7 @@ class PostgreSQLStore:
|
|
|
password=os.getenv('KNOWHUB_PASSWORD'),
|
|
|
database=os.getenv('KNOWHUB_DB_NAME')
|
|
|
)
|
|
|
- self.conn.autocommit = False
|
|
|
+ self.conn.autocommit = True
|
|
|
|
|
|
def _ensure_connection(self):
|
|
|
if self.conn.closed != 0:
|
|
|
@@ -113,18 +148,24 @@ class PostgreSQLStore:
|
|
|
))
|
|
|
# 写入关联表
|
|
|
kid = knowledge['id']
|
|
|
- for req_id in knowledge.get('requirement_ids', []):
|
|
|
+ req_links = _normalize_links(knowledge, 'requirement_links', 'requirement_ids', 'related') or []
|
|
|
+ for req_id, rtype in req_links:
|
|
|
cursor.execute(
|
|
|
- "INSERT INTO requirement_knowledge (requirement_id, knowledge_id) VALUES (%s, %s) ON CONFLICT DO NOTHING",
|
|
|
- (req_id, kid))
|
|
|
- for cap_id in knowledge.get('capability_ids', []):
|
|
|
+ "INSERT INTO requirement_knowledge (requirement_id, knowledge_id, relation_type) "
|
|
|
+ "VALUES (%s, %s, %s) ON CONFLICT DO NOTHING",
|
|
|
+ (req_id, kid, rtype))
|
|
|
+ cap_links = _normalize_links(knowledge, 'capability_links', 'capability_ids', 'related') or []
|
|
|
+ for cap_id, rtype in cap_links:
|
|
|
cursor.execute(
|
|
|
- "INSERT INTO capability_knowledge (capability_id, knowledge_id) VALUES (%s, %s) ON CONFLICT DO NOTHING",
|
|
|
- (cap_id, kid))
|
|
|
- for tool_id in knowledge.get('tool_ids', []):
|
|
|
+ "INSERT INTO capability_knowledge (capability_id, knowledge_id, relation_type) "
|
|
|
+ "VALUES (%s, %s, %s) ON CONFLICT DO NOTHING",
|
|
|
+ (cap_id, kid, rtype))
|
|
|
+ tool_links = _normalize_links(knowledge, 'tool_links', 'tool_ids', 'related') or []
|
|
|
+ for tool_id, rtype in tool_links:
|
|
|
cursor.execute(
|
|
|
- "INSERT INTO tool_knowledge (tool_id, knowledge_id) VALUES (%s, %s) ON CONFLICT DO NOTHING",
|
|
|
- (tool_id, kid))
|
|
|
+ "INSERT INTO tool_knowledge (tool_id, knowledge_id, relation_type) "
|
|
|
+ "VALUES (%s, %s, %s) ON CONFLICT DO NOTHING",
|
|
|
+ (tool_id, kid, rtype))
|
|
|
for res_id in knowledge.get('resource_ids', []):
|
|
|
cursor.execute(
|
|
|
"INSERT INTO knowledge_resource (knowledge_id, resource_id) VALUES (%s, %s) ON CONFLICT DO NOTHING",
|
|
|
@@ -220,10 +261,10 @@ class PostgreSQLStore:
|
|
|
cursor = self._get_cursor()
|
|
|
try:
|
|
|
# 分离关联字段和实体字段
|
|
|
- req_ids = updates.pop('requirement_ids', None)
|
|
|
- cap_ids = updates.pop('capability_ids', None)
|
|
|
- tool_ids = updates.pop('tool_ids', None)
|
|
|
- resource_ids = updates.pop('resource_ids', None)
|
|
|
+ rel_keys = ('requirement_ids', 'requirement_links',
|
|
|
+ 'capability_ids', 'capability_links',
|
|
|
+ 'tool_ids', 'tool_links', 'resource_ids')
|
|
|
+ rel_data = {k: updates.pop(k) for k in rel_keys if k in updates}
|
|
|
|
|
|
if updates:
|
|
|
set_parts = []
|
|
|
@@ -240,30 +281,36 @@ class PostgreSQLStore:
|
|
|
cursor.execute(sql, params)
|
|
|
|
|
|
# 更新关联表(全量替换)
|
|
|
- if req_ids is not None:
|
|
|
+ req_links = _normalize_links(rel_data, 'requirement_links', 'requirement_ids', 'related')
|
|
|
+ if req_links is not None:
|
|
|
cursor.execute("DELETE FROM requirement_knowledge WHERE knowledge_id = %s", (knowledge_id,))
|
|
|
- for req_id in req_ids:
|
|
|
+ for req_id, rtype in req_links:
|
|
|
cursor.execute(
|
|
|
- "INSERT INTO requirement_knowledge (requirement_id, knowledge_id) VALUES (%s, %s) ON CONFLICT DO NOTHING",
|
|
|
- (req_id, knowledge_id))
|
|
|
+ "INSERT INTO requirement_knowledge (requirement_id, knowledge_id, relation_type) "
|
|
|
+ "VALUES (%s, %s, %s) ON CONFLICT DO NOTHING",
|
|
|
+ (req_id, knowledge_id, rtype))
|
|
|
|
|
|
- if cap_ids is not None:
|
|
|
+ cap_links = _normalize_links(rel_data, 'capability_links', 'capability_ids', 'related')
|
|
|
+ if cap_links is not None:
|
|
|
cursor.execute("DELETE FROM capability_knowledge WHERE knowledge_id = %s", (knowledge_id,))
|
|
|
- for cap_id in cap_ids:
|
|
|
+ for cap_id, rtype in cap_links:
|
|
|
cursor.execute(
|
|
|
- "INSERT INTO capability_knowledge (capability_id, knowledge_id) VALUES (%s, %s) ON CONFLICT DO NOTHING",
|
|
|
- (cap_id, knowledge_id))
|
|
|
+ "INSERT INTO capability_knowledge (capability_id, knowledge_id, relation_type) "
|
|
|
+ "VALUES (%s, %s, %s) ON CONFLICT DO NOTHING",
|
|
|
+ (cap_id, knowledge_id, rtype))
|
|
|
|
|
|
- if tool_ids is not None:
|
|
|
+ tool_links = _normalize_links(rel_data, 'tool_links', 'tool_ids', 'related')
|
|
|
+ if tool_links is not None:
|
|
|
cursor.execute("DELETE FROM tool_knowledge WHERE knowledge_id = %s", (knowledge_id,))
|
|
|
- for tool_id in tool_ids:
|
|
|
+ for tool_id, rtype in tool_links:
|
|
|
cursor.execute(
|
|
|
- "INSERT INTO tool_knowledge (tool_id, knowledge_id) VALUES (%s, %s) ON CONFLICT DO NOTHING",
|
|
|
- (tool_id, knowledge_id))
|
|
|
+ "INSERT INTO tool_knowledge (tool_id, knowledge_id, relation_type) "
|
|
|
+ "VALUES (%s, %s, %s) ON CONFLICT DO NOTHING",
|
|
|
+ (tool_id, knowledge_id, rtype))
|
|
|
|
|
|
- if resource_ids is not None:
|
|
|
+ if 'resource_ids' in rel_data and rel_data['resource_ids'] is not None:
|
|
|
cursor.execute("DELETE FROM knowledge_resource WHERE knowledge_id = %s", (knowledge_id,))
|
|
|
- for res_id in resource_ids:
|
|
|
+ for res_id in rel_data['resource_ids']:
|
|
|
cursor.execute(
|
|
|
"INSERT INTO knowledge_resource (knowledge_id, resource_id) VALUES (%s, %s) ON CONFLICT DO NOTHING",
|
|
|
(knowledge_id, res_id))
|
|
|
@@ -303,6 +350,45 @@ class PostgreSQLStore:
|
|
|
finally:
|
|
|
cursor.close()
|
|
|
|
|
|
+ def add_requirement(self, knowledge_id: str, requirement_id: str,
|
|
|
+ relation_type: str = 'related'):
|
|
|
+ """增量挂接 requirement-knowledge 边"""
|
|
|
+ cursor = self._get_cursor()
|
|
|
+ try:
|
|
|
+ cursor.execute(
|
|
|
+ "INSERT INTO requirement_knowledge (requirement_id, knowledge_id, relation_type) "
|
|
|
+ "VALUES (%s, %s, %s) ON CONFLICT DO NOTHING",
|
|
|
+ (requirement_id, knowledge_id, relation_type))
|
|
|
+ self.conn.commit()
|
|
|
+ finally:
|
|
|
+ cursor.close()
|
|
|
+
|
|
|
+ def add_capability(self, knowledge_id: str, capability_id: str,
|
|
|
+ relation_type: str = 'related'):
|
|
|
+ """增量挂接 capability-knowledge 边"""
|
|
|
+ cursor = self._get_cursor()
|
|
|
+ try:
|
|
|
+ cursor.execute(
|
|
|
+ "INSERT INTO capability_knowledge (capability_id, knowledge_id, relation_type) "
|
|
|
+ "VALUES (%s, %s, %s) ON CONFLICT DO NOTHING",
|
|
|
+ (capability_id, knowledge_id, relation_type))
|
|
|
+ self.conn.commit()
|
|
|
+ finally:
|
|
|
+ cursor.close()
|
|
|
+
|
|
|
+ def add_tool(self, knowledge_id: str, tool_id: str,
|
|
|
+ relation_type: str = 'related'):
|
|
|
+ """增量挂接 tool-knowledge 边"""
|
|
|
+ cursor = self._get_cursor()
|
|
|
+ try:
|
|
|
+ cursor.execute(
|
|
|
+ "INSERT INTO tool_knowledge (tool_id, knowledge_id, relation_type) "
|
|
|
+ "VALUES (%s, %s, %s) ON CONFLICT DO NOTHING",
|
|
|
+ (tool_id, knowledge_id, relation_type))
|
|
|
+ self.conn.commit()
|
|
|
+ finally:
|
|
|
+ cursor.close()
|
|
|
+
|
|
|
def count(self) -> int:
|
|
|
"""返回知识总数"""
|
|
|
cursor = self._get_cursor()
|
|
|
@@ -349,7 +435,8 @@ class PostgreSQLStore:
|
|
|
if 'eval' in result and isinstance(result['eval'], str):
|
|
|
result['eval'] = json.loads(result['eval'])
|
|
|
# 关联字段(来自 junction table 子查询,可能是 JSON 字符串或已解析的列表)
|
|
|
- for field in ('requirement_ids', 'capability_ids', 'tool_ids', 'resource_ids'):
|
|
|
+ for field in ('requirement_ids', 'capability_ids', 'tool_ids', 'resource_ids',
|
|
|
+ 'requirement_links', 'capability_links', 'tool_links'):
|
|
|
if field in result and isinstance(result[field], str):
|
|
|
result[field] = json.loads(result[field])
|
|
|
elif field in result and result[field] is None:
|
|
|
@@ -400,18 +487,24 @@ class PostgreSQLStore:
|
|
|
# 批量写入关联表
|
|
|
for k in knowledge_list:
|
|
|
kid = k['id']
|
|
|
- for req_id in k.get('requirement_ids', []):
|
|
|
+ req_links = _normalize_links(k, 'requirement_links', 'requirement_ids', 'related') or []
|
|
|
+ for req_id, rtype in req_links:
|
|
|
cursor.execute(
|
|
|
- "INSERT INTO requirement_knowledge (requirement_id, knowledge_id) VALUES (%s, %s) ON CONFLICT DO NOTHING",
|
|
|
- (req_id, kid))
|
|
|
- for cap_id in k.get('capability_ids', []):
|
|
|
+ "INSERT INTO requirement_knowledge (requirement_id, knowledge_id, relation_type) "
|
|
|
+ "VALUES (%s, %s, %s) ON CONFLICT DO NOTHING",
|
|
|
+ (req_id, kid, rtype))
|
|
|
+ cap_links = _normalize_links(k, 'capability_links', 'capability_ids', 'related') or []
|
|
|
+ for cap_id, rtype in cap_links:
|
|
|
cursor.execute(
|
|
|
- "INSERT INTO capability_knowledge (capability_id, knowledge_id) VALUES (%s, %s) ON CONFLICT DO NOTHING",
|
|
|
- (cap_id, kid))
|
|
|
- for tool_id in k.get('tool_ids', []):
|
|
|
+ "INSERT INTO capability_knowledge (capability_id, knowledge_id, relation_type) "
|
|
|
+ "VALUES (%s, %s, %s) ON CONFLICT DO NOTHING",
|
|
|
+ (cap_id, kid, rtype))
|
|
|
+ tool_links = _normalize_links(k, 'tool_links', 'tool_ids', 'related') or []
|
|
|
+ for tool_id, rtype in tool_links:
|
|
|
cursor.execute(
|
|
|
- "INSERT INTO tool_knowledge (tool_id, knowledge_id) VALUES (%s, %s) ON CONFLICT DO NOTHING",
|
|
|
- (tool_id, kid))
|
|
|
+ "INSERT INTO tool_knowledge (tool_id, knowledge_id, relation_type) "
|
|
|
+ "VALUES (%s, %s, %s) ON CONFLICT DO NOTHING",
|
|
|
+ (tool_id, kid, rtype))
|
|
|
for res_id in k.get('resource_ids', []):
|
|
|
cursor.execute(
|
|
|
"INSERT INTO knowledge_resource (knowledge_id, resource_id) VALUES (%s, %s) ON CONFLICT DO NOTHING",
|