|
@@ -27,7 +27,29 @@ class PostgreSQLToolStore:
|
|
|
self.conn.autocommit = False
|
|
self.conn.autocommit = False
|
|
|
print(f"[PostgreSQL Tool] 已连接到远程数据库: {os.getenv('KNOWHUB_DB')}")
|
|
print(f"[PostgreSQL Tool] 已连接到远程数据库: {os.getenv('KNOWHUB_DB')}")
|
|
|
|
|
|
|
|
|
|
+ def _reconnect(self):
|
|
|
|
|
+ self.conn = psycopg2.connect(
|
|
|
|
|
+ host=os.getenv('KNOWHUB_DB'),
|
|
|
|
|
+ port=int(os.getenv('KNOWHUB_PORT', 5432)),
|
|
|
|
|
+ user=os.getenv('KNOWHUB_USER'),
|
|
|
|
|
+ password=os.getenv('KNOWHUB_PASSWORD'),
|
|
|
|
|
+ database=os.getenv('KNOWHUB_DB_NAME')
|
|
|
|
|
+ )
|
|
|
|
|
+ self.conn.autocommit = False
|
|
|
|
|
+
|
|
|
|
|
+ def _ensure_connection(self):
|
|
|
|
|
+ if self.conn.closed != 0:
|
|
|
|
|
+ self._reconnect()
|
|
|
|
|
+ else:
|
|
|
|
|
+ try:
|
|
|
|
|
+ c = self.conn.cursor()
|
|
|
|
|
+ c.execute("SELECT 1")
|
|
|
|
|
+ c.close()
|
|
|
|
|
+ except (psycopg2.OperationalError, psycopg2.InterfaceError):
|
|
|
|
|
+ self._reconnect()
|
|
|
|
|
+
|
|
|
def _get_cursor(self):
|
|
def _get_cursor(self):
|
|
|
|
|
+ self._ensure_connection()
|
|
|
return self.conn.cursor(cursor_factory=RealDictCursor)
|
|
return self.conn.cursor(cursor_factory=RealDictCursor)
|
|
|
|
|
|
|
|
def insert_or_update(self, tool: Dict):
|
|
def insert_or_update(self, tool: Dict):
|
|
@@ -38,8 +60,8 @@ class PostgreSQLToolStore:
|
|
|
INSERT INTO tool_table (
|
|
INSERT INTO tool_table (
|
|
|
id, name, version, introduction, tutorial, input, output,
|
|
id, name, version, introduction, tutorial, input, output,
|
|
|
updated_time, status, capabilities, tool_knowledge,
|
|
updated_time, status, capabilities, tool_knowledge,
|
|
|
- case_knowledge, process_knowledge, embedding
|
|
|
|
|
- ) VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s)
|
|
|
|
|
|
|
+ case_knowledge, process_knowledge, embedding, implemented_tool_ids
|
|
|
|
|
+ ) VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s)
|
|
|
ON CONFLICT (id) DO UPDATE SET
|
|
ON CONFLICT (id) DO UPDATE SET
|
|
|
name = EXCLUDED.name,
|
|
name = EXCLUDED.name,
|
|
|
version = EXCLUDED.version,
|
|
version = EXCLUDED.version,
|
|
@@ -53,7 +75,8 @@ class PostgreSQLToolStore:
|
|
|
tool_knowledge = EXCLUDED.tool_knowledge,
|
|
tool_knowledge = EXCLUDED.tool_knowledge,
|
|
|
case_knowledge = EXCLUDED.case_knowledge,
|
|
case_knowledge = EXCLUDED.case_knowledge,
|
|
|
process_knowledge = EXCLUDED.process_knowledge,
|
|
process_knowledge = EXCLUDED.process_knowledge,
|
|
|
- embedding = EXCLUDED.embedding
|
|
|
|
|
|
|
+ embedding = EXCLUDED.embedding,
|
|
|
|
|
+ implemented_tool_ids = EXCLUDED.implemented_tool_ids
|
|
|
""", (
|
|
""", (
|
|
|
tool['id'],
|
|
tool['id'],
|
|
|
tool.get('name', ''),
|
|
tool.get('name', ''),
|
|
@@ -69,6 +92,7 @@ class PostgreSQLToolStore:
|
|
|
json.dumps(tool.get('case_knowledge', [])),
|
|
json.dumps(tool.get('case_knowledge', [])),
|
|
|
json.dumps(tool.get('process_knowledge', [])),
|
|
json.dumps(tool.get('process_knowledge', [])),
|
|
|
tool.get('embedding'),
|
|
tool.get('embedding'),
|
|
|
|
|
+ json.dumps(tool.get('implemented_tool_ids', [])),
|
|
|
))
|
|
))
|
|
|
self.conn.commit()
|
|
self.conn.commit()
|
|
|
finally:
|
|
finally:
|
|
@@ -81,7 +105,7 @@ class PostgreSQLToolStore:
|
|
|
cursor.execute("""
|
|
cursor.execute("""
|
|
|
SELECT id, name, version, introduction, tutorial, input, output,
|
|
SELECT id, name, version, introduction, tutorial, input, output,
|
|
|
updated_time, status, capabilities, tool_knowledge,
|
|
updated_time, status, capabilities, tool_knowledge,
|
|
|
- case_knowledge, process_knowledge
|
|
|
|
|
|
|
+ case_knowledge, process_knowledge, implemented_tool_ids
|
|
|
FROM tool_table WHERE id = %s
|
|
FROM tool_table WHERE id = %s
|
|
|
""", (tool_id,))
|
|
""", (tool_id,))
|
|
|
result = cursor.fetchone()
|
|
result = cursor.fetchone()
|
|
@@ -101,7 +125,7 @@ class PostgreSQLToolStore:
|
|
|
sql = f"""
|
|
sql = f"""
|
|
|
SELECT id, name, version, introduction, tutorial, input, output,
|
|
SELECT id, name, version, introduction, tutorial, input, output,
|
|
|
updated_time, status, capabilities, tool_knowledge,
|
|
updated_time, status, capabilities, tool_knowledge,
|
|
|
- case_knowledge, process_knowledge,
|
|
|
|
|
|
|
+ case_knowledge, process_knowledge, implemented_tool_ids,
|
|
|
1 - (embedding <=> %s::real[]) as score
|
|
1 - (embedding <=> %s::real[]) as score
|
|
|
FROM tool_table
|
|
FROM tool_table
|
|
|
WHERE embedding IS NOT NULL AND status = %s
|
|
WHERE embedding IS NOT NULL AND status = %s
|
|
@@ -112,7 +136,7 @@ class PostgreSQLToolStore:
|
|
|
sql = f"""
|
|
sql = f"""
|
|
|
SELECT id, name, version, introduction, tutorial, input, output,
|
|
SELECT id, name, version, introduction, tutorial, input, output,
|
|
|
updated_time, status, capabilities, tool_knowledge,
|
|
updated_time, status, capabilities, tool_knowledge,
|
|
|
- case_knowledge, process_knowledge,
|
|
|
|
|
|
|
+ case_knowledge, process_knowledge, implemented_tool_ids,
|
|
|
1 - (embedding <=> %s::real[]) as score
|
|
1 - (embedding <=> %s::real[]) as score
|
|
|
FROM tool_table
|
|
FROM tool_table
|
|
|
WHERE embedding IS NOT NULL
|
|
WHERE embedding IS NOT NULL
|
|
@@ -133,7 +157,7 @@ class PostgreSQLToolStore:
|
|
|
cursor.execute("""
|
|
cursor.execute("""
|
|
|
SELECT id, name, version, introduction, tutorial, input, output,
|
|
SELECT id, name, version, introduction, tutorial, input, output,
|
|
|
updated_time, status, capabilities, tool_knowledge,
|
|
updated_time, status, capabilities, tool_knowledge,
|
|
|
- case_knowledge, process_knowledge
|
|
|
|
|
|
|
+ case_knowledge, process_knowledge, implemented_tool_ids
|
|
|
FROM tool_table
|
|
FROM tool_table
|
|
|
WHERE status = %s
|
|
WHERE status = %s
|
|
|
ORDER BY updated_time DESC
|
|
ORDER BY updated_time DESC
|
|
@@ -143,7 +167,7 @@ class PostgreSQLToolStore:
|
|
|
cursor.execute("""
|
|
cursor.execute("""
|
|
|
SELECT id, name, version, introduction, tutorial, input, output,
|
|
SELECT id, name, version, introduction, tutorial, input, output,
|
|
|
updated_time, status, capabilities, tool_knowledge,
|
|
updated_time, status, capabilities, tool_knowledge,
|
|
|
- case_knowledge, process_knowledge
|
|
|
|
|
|
|
+ case_knowledge, process_knowledge, implemented_tool_ids
|
|
|
FROM tool_table
|
|
FROM tool_table
|
|
|
ORDER BY updated_time DESC
|
|
ORDER BY updated_time DESC
|
|
|
LIMIT %s OFFSET %s
|
|
LIMIT %s OFFSET %s
|
|
@@ -203,7 +227,7 @@ class PostgreSQLToolStore:
|
|
|
return None
|
|
return None
|
|
|
result = dict(row)
|
|
result = dict(row)
|
|
|
for field in ('input', 'output', 'capabilities', 'tool_knowledge',
|
|
for field in ('input', 'output', 'capabilities', 'tool_knowledge',
|
|
|
- 'case_knowledge', 'process_knowledge'):
|
|
|
|
|
|
|
+ 'case_knowledge', 'process_knowledge', 'implemented_tool_ids'):
|
|
|
if field in result and isinstance(result[field], str):
|
|
if field in result and isinstance(result[field], str):
|
|
|
try:
|
|
try:
|
|
|
result[field] = json.loads(result[field]) if result[field].strip() else None
|
|
result[field] = json.loads(result[field]) if result[field].strip() else None
|