From 9bb26f0704f3340353225670d487e2e3c39c103e Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E5=8F=B2=E9=B8=BF=E5=AE=87?= Date: Tue, 18 Feb 2025 15:23:53 +0800 Subject: [PATCH 1/2] =?UTF-8?q?mock:=20=E6=96=B0=E5=A2=9E=20A-Ops=20Apollo?= =?UTF-8?q?=20=E7=9B=B8=E5=85=B3=E6=8E=A5=E5=8F=A3=E6=95=B0=E6=8D=AE?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: 史鸿宇 --- mock/make_data.py | 654 +++++++++++++++++++++++++++++++--------------- 1 file changed, 446 insertions(+), 208 deletions(-) diff --git a/mock/make_data.py b/mock/make_data.py index 5d45a5742..7d8ee314a 100644 --- a/mock/make_data.py +++ b/mock/make_data.py @@ -1,27 +1,25 @@ -from typing import Any -import uuid import asyncio +import urllib.parse +import uuid from datetime import datetime, timezone -from typing import Optional, List -from pydantic import BaseModel, Field from enum import Enum +from typing import Any, Optional + +from pydantic import BaseModel, Field from pymongo import MongoClient from pymongo.errors import PyMongoError -import urllib.parse -from bson import ObjectId -from pydantic import BaseModel, Field, HttpUrl -# 假设 PermissionType 已经在别处定义 +# 假设 PermissionType 已经在别处定义 class PermissionType(str, Enum): PRIVATE = "private" PUBLIC = "public" -# 定义模型类 +# 定义模型类 class PoolBase(BaseModel): - id: str = Field(alias='_id') + id: str = Field(alias="_id") name: str description: str created_at: float = Field(default_factory=lambda: round(datetime.now(tz=timezone.utc).timestamp(), 3)) @@ -56,11 +54,11 @@ class ServicePool(PoolBase): # } # MongoDB配置 config = { - 'MONGODB_USER': 'euler_copilot', - 'MONGODB_PWD': '8URM%HtCHQPxKe$u', - 'MONGODB_HOST': '10.43.208.180', - 'MONGODB_PORT': '27017', - 'MONGODB_DATABASE': 'euler_copilot' + "MONGODB_USER": "euler_copilot", + "MONGODB_PWD": "8URM%HtCHQPxKe$u", + "MONGODB_HOST": "10.43.208.180", + "MONGODB_PORT": "27017", + "MONGODB_DATABASE": "euler_copilot", } @@ -83,44 +81,42 @@ async def insert_service_pool(): api_info_1 = ServiceApiInfo(filename="example_1.yaml", description="Example API 1", path="/api/example/3") api_info_2 = ServiceApiInfo(filename="example_2.yaml", description="Example API 2", path="/api/example/2") api_info_3 = ServiceApiInfo(filename="example_3.yaml", description="Example API 3", path="/api/example/1") - id = "6a08c845-abdc-45fb-853e-54a806437dab" - print(id) + sys_id = "6a08c845-abdc-45fb-853e-54a806437dab" service_pool_sys = ServicePool( - _id=id, + _id=sys_id, name="系统", description="系统函数", author="test", api=[api_info_1, api_info_2, api_info_3], - permissions=Permission(type=PermissionType.PUBLIC.value, users=["user1", "user2"]).dict(), - favorites=["user1","test"], - hashes={"file1": "hash1", "file2": "hash2"} + permissions=Permission(type=PermissionType.PUBLIC, users=["user1", "user2"]), + favorites=["user1", "test"], + hashes={"file1": "hash1", "file2": "hash2"}, ) - id = str(uuid.uuid4()) - print(id) + aops_id = "1137ab09-20ae-4278-8346-524d4ce81d2f" service_pool_a_ops = ServicePool( - _id=id, - name="a-ops-apollo", + _id=aops_id, + name="aops-apollo", description="a-ops下cve相关组件", author="test", api=[api_info_1, api_info_2, api_info_3], - permissions=Permission(type=PermissionType.PUBLIC.value, users=["user1", "user2"]).dict(), + permissions=Permission(type=PermissionType.PUBLIC, users=["user1", "user2"]), favorites=["user1"], - hashes={"file1": "hash1", "file2": "hash2"} + hashes={"file1": "hash1", "file2": "hash2"}, ) """插入ServicePool实例到MongoDB""" - collection = MongoDB.get_collection('service') + collection = MongoDB.get_collection("service") result = collection.delete_many({}) # 将Pydantic模型转换为字典并插入到MongoDB中 try: result = collection.update_one( - {'_id': service_pool_sys.id}, # 查找条件 + {"_id": service_pool_sys.id}, # 查找条件 {"$set": service_pool_sys.dict(by_alias=True)}, # 更新操作 - upsert=True # 如果不存在则插入新文档 + upsert=True, # 如果不存在则插入新文档 ) result = collection.update_one( - {'_id': service_pool_a_ops.id}, # 查找条件 + {"_id": service_pool_a_ops.id}, # 查找条件 {"$set": service_pool_a_ops.dict(by_alias=True)}, # 更新操作 - upsert=True # 如果不存在则插入新文档 + upsert=True, # 如果不存在则插入新文档 ) print(f"Inserted document with id: {result.upserted_id}") except PyMongoError as e: @@ -146,43 +142,48 @@ class NodePool(PoolBase): output_schema: dict[str, Any] = Field(description="Node的输出schema；做输出的展示用", default={}) -async def insert_node_pool(): - collection = MongoDB.get_collection('node') +async def insert_node_pool() -> None: + collection = MongoDB.get_collection("node") result = collection.delete_many({}) # 清空集合中的所有文档（仅用于演示） - node_pools = [NodePool( - _id=str(uuid.uuid4()), # 自动生成一个唯一的 ID - service_id="6a08c845-abdc-45fb-853e-54a806437dab", # 使用 "test" 作为 service_id - call_id="knowledge_base", # 随机生成一个 call_id - name="【KNOWLEDGE】知识库", # 提供名称 - description="支持知识库中文档的查询", # 提供描述 - params_schema={ - "search_methods": [], - "rerank_methods": [], - "konwledge_base_id": "", - "query": "", - "top_k": 0}, - output_schema={"content": {"type":"string", "description": "回答"}} - ), NodePool( - _id=str(uuid.uuid4()), # 自动生成一个唯一的 ID - service_id="6a08c845-abdc-45fb-853e-54a806437dab", # 使用 "test" 作为 service_id - call_id="LLM", # 随机生成一个 call_id - name="【LLM】大模型", # 提供名称 - description="大模型调用", # 提供描述 - params_schema={ - "base_url": "", - "api_key": "", - "max_tokens": 0, - "is_stream": True, - "prompt": "", - "temperature": 0}, - output_schema={"content": {}} - ), NodePool( - _id=str(uuid.uuid4()), # 自动生成一个唯一的 ID - service_id="6a08c845-abdc-45fb-853e-54a806437dab", # 使用 "test" 作为 service_id - call_id="choice", # 随机生成一个 call_id - name="【LLM】意图识别", # 提供名称 - description="利用大模型能力选择分支", # 提供描述 - params_schema={ + node_pools = [ + NodePool( + _id=str(uuid.uuid4()), # 自动生成一个唯一的 ID + service_id="6a08c845-abdc-45fb-853e-54a806437dab", # 使用 "test" 作为 service_id + call_id="knowledge_base", # 随机生成一个 call_id + name="【KNOWLEDGE】知识库", # 提供名称 + description="支持知识库中文档的查询", # 提供描述 + params_schema={ + "search_methods": [], + "rerank_methods": [], + "konwledge_base_id": "", + "query": "", + "top_k": 0, + }, + output_schema={"content": {"type": "string", "description": "回答"}}, + ), + NodePool( + _id=str(uuid.uuid4()), # 自动生成一个唯一的 ID + service_id="6a08c845-abdc-45fb-853e-54a806437dab", # 使用 "test" 作为 service_id + call_id="LLM", # 随机生成一个 call_id + name="【LLM】大模型", # 提供名称 + description="大模型调用", # 提供描述 + params_schema={ + "base_url": "", + "api_key": "", + "max_tokens": 0, + "is_stream": True, + "prompt": "", + "temperature": 0, + }, + output_schema={"content": {}}, + ), + NodePool( + _id=str(uuid.uuid4()), # 自动生成一个唯一的 ID + service_id="6a08c845-abdc-45fb-853e-54a806437dab", # 使用 "test" 作为 service_id + call_id="choice", # 随机生成一个 call_id + name="【LLM】意图识别", # 提供名称 + description="利用大模型能力选择分支", # 提供描述 + params_schema={ "choices": [ { "branchId": "source_a", @@ -195,132 +196,362 @@ async def insert_node_pool(): "description": "ELSE B", }, ], - }, - output_schema={ - "type":"object", - "properties": { - "content": { - "type":"string", - "description": "回答", + }, + output_schema={ + "type": "object", + "properties": { + "content": { + "type": "string", + "description": "回答", + }, }, }, - } - ), NodePool( - _id=str(uuid.uuid4()), # 自动生成一个唯一的 ID - service_id="6a08c845-abdc-45fb-853e-54a806437dab", # 使用 "test" 作为 service_id - call_id="choice", # 随机生成一个 call_id - name="【CHOICE】条件分支", # 提供名称 - description="条件分支节点", # 提供描述 - params_schema={ - "choices": [ - { - "branchId":"source_a", - "description":"IF A", - "operator":"", - "variable_a":"", - "variable_b":"", + ), + NodePool( + _id=str(uuid.uuid4()), # 自动生成一个唯一的 ID + service_id="6a08c845-abdc-45fb-853e-54a806437dab", # 使用 "test" 作为 service_id + call_id="choice", # 随机生成一个 call_id + name="【CHOICE】条件分支", # 提供名称 + description="条件分支节点", # 提供描述 + params_schema={ + "choices": [ + { + "branchId": "source_a", + "description": "IF A", + "operator": "", + "variable_a": "", + "variable_b": "", + }, + { + "branchId": "source_b", + "description": "ELSE B", + }, + ], + }, + output_schema={}, + ), + NodePool( + _id=str(uuid.uuid4()), # 自动生成一个唯一的 ID + service_id="6a08c845-abdc-45fb-853e-54a806437dab", # 使用 "test" 作为 service_id + call_id="loop_begin", # 随机生成一个 call_id + name="【LOOP】循环开始节点", # 提供名称 + description="", # 提供描述 + params_schema={"operation_exp": {}}, + output_schema={}, + ), + NodePool( + _id=str(uuid.uuid4()), # 自动生成一个唯一的 ID + service_id="6a08c845-abdc-45fb-853e-54a806437dab", # 使用 "test" 作为 service_id + call_id="loop_begin", # 随机生成一个 call_id + name="【LOOP】循环结束节点", # 提供名称 + description="", # 提供描述 + params_schema={"operation_exp": {}}, + output_schema={}, + ), + NodePool( + _id=str(uuid.uuid4()), # 自动生成一个唯一的 ID + service_id="6a08c845-abdc-45fb-853e-54a806437dab", # 使用 "test" 作为 service_id + call_id="template_exchange", # 随机生成一个 call_id + name="【LLM】模板转换", # 提供名称 + description="This is an example node pool for demonstration purposes.", # 提供描述 + params_schema={"input_schema": {}, "exchange_rule": [{}]}, + output_schema={ + "type": "object", + "properties": { + "output_schema": { + "type": "object", + "description": "嵌套字典结构", + "properties": { + "content": { + "type": "string", + "description": "回答", }, - { - "branchId":"source_b", - "description":"ELSE B", + "task_id": { + "type": "string", + "description": "任务ID", }, - ], + }, + }, + }, + }, + ), + NodePool( + _id=str(uuid.uuid4()), # 自动生成一个唯一的 ID + service_id="1137ab09-20ae-4278-8346-524d4ce81d2f", + call_id="api", + name="【API】获取任务简介", + description="调用接口，获取已知的任务列表与任务的基本信息（名称、创建时间、任务类型等）", + params_schema={ + "full_url": "https://a-ops3.local/vulnerabilities/task/list/get", + "service_id": "aops-apollo", + "method": "post", + "input_data": { + "page": 1, + "page_size": 10, + "filter": { + "cluster_list": [], + }, }, - output_schema={} - ), NodePool( - _id=str(uuid.uuid4()), # 自动生成一个唯一的 ID - service_id="6a08c845-abdc-45fb-853e-54a806437dab", # 使用 "test" 作为 service_id - call_id="loop_begin", # 随机生成一个 call_id - name="【LOOP】循环开始节点", # 提供名称 - description="", # 提供描述 - params_schema={"operation_exp": {}}, - output_schema={}, - ), NodePool( - _id=str(uuid.uuid4()), # 自动生成一个唯一的 ID - service_id="6a08c845-abdc-45fb-853e-54a806437dab", # 使用 "test" 作为 service_id - call_id="loop_begin", # 随机生成一个 call_id - name="【LOOP】循环结束节点", # 提供名称 - description="", # 提供描述 - params_schema={"operation_exp": {}}, - output_schema={}, - ), NodePool( - _id=str(uuid.uuid4()), # 自动生成一个唯一的 ID - service_id="6a08c845-abdc-45fb-853e-54a806437dab", # 使用 "test" 作为 service_id - call_id="template_exchange", # 随机生成一个 call_id - name="【LLM】模板转换", # 提供名称 - description="This is an example node pool for demonstration purposes.", # 提供描述 - params_schema={"input_schema": {}, - "exchange_rule": [{}]}, - output_schema={ - "type":"object", - "properties": { - "output_schema": { - "type":"object", - "description": "嵌套字典结构", - "properties": { - "content": { - "type":"string", - "description": "回答", + "timeout": 300, + "output_key": [ + { + "key": "data.result", + "path": "task_list", + }, + ], + }, + output_schema={ + "type": "object", + "properties": { + "status_code": { + "type": "integer", + "description": "HTTP状态码", + }, + "data": { + "type": "object", + "description": "接口的返回数据", + "properties": { + "task_list": { + "type": "array", + "description": "任务列表", + "items": { + "type": "object", + "properties": { + "task_id": { + "type": "string", + "description": "任务ID", + }, + "task_name": { + "type": "string", + "description": "任务名称", + }, + "task_type": { + "type": "string", + "description": "任务类型", + "enum": ["cve_scan", "cve_fix"], + }, + }, + }, + }, }, - "task_id": { - "type":"string", - "description": "任务ID", - } }, }, }, - }, - ), NodePool( - _id=str(uuid.uuid4()), # 自动生成一个唯一的 ID - service_id="6d7f65ff-55de-4206-9e5b-87aa2b529bd4", # 使用 "test" 作为 service_id - call_id="search", # 随机生成一个 call_id - name="【API】扫描CVE漏洞", # 提供名称 - description="扫描某个机器所有的CVE漏洞", # 提供描述 - params_schema={"host": ""}, - output_schema={"task_id": ""}, - ), NodePool( - _id=str(uuid.uuid4()), # 自动生成一个唯一的 ID - service_id="6d7f65ff-55de-4206-9e5b-87aa2b529bd4", # 使用 "test" 作为 service_id - call_id="search", # 随机生成一个 call_id - name="【API】修复CVE漏洞", # 提供名称 - description="修复某个机器上的CVE漏洞", # 提供描述 - params_schema={"host": "", - "cve_id": ""}, - output_schema={"task_report":{}}, - ) + ), + NodePool( + _id=str(uuid.uuid4()), # 自动生成一个唯一的 ID + service_id="1137ab09-20ae-4278-8346-524d4ce81d2f", + call_id="choice", + name="【CHOICE】判断任务类型", + description="调用意图识别工具，判断任务列表中的第一个任务的类型", + params_schema={ + "choices": [ + { + "branchId": "is_scan", + "description": '任务类型为"CVE修复任务"', + "propose": '当值为cve_scan时，任务类型为"CVE修复任务"，选择此分支', + "variable_a": "{{input.task_list[0].task_type}}", + }, + { + "branchId": "is_fix", + "description": '任务类型为"CVE修复任务"', + "propose": '当值为cve_fix时，任务类型为"CVE修复任务"，选择此分支', + "variable_a": "{{input.task_list[0].task_type}}", + }, + ], + }, + output_schema={ + "type": "object", + "properties": {}, + }, + ), + NodePool( + _id=str(uuid.uuid4()), # 自动生成一个唯一的 ID + service_id="1137ab09-20ae-4278-8346-524d4ce81d2f", + call_id="api", + name="【API】获取CVE修复的结果", + description="调用接口，获取特定CVE修复任务的最终结果", + params_schema={ + "full_url": "https://a-ops3.local/vulnerabilities/task/cve_fix/result/get", + "service_id": "aops-apollo", + "method": "post", + "input_data": {}, + "timeout": 300, + "output_key": [ + { + "key": "data", + "path": "result", + }, + ], + }, + output_schema={ + "type": "object", + "description": "API的返回信息", + "properties": { + "status_code": {"type": "integer", "description": "HTTP状态码"}, + "data": { + "type": "object", + "description": "接口的返回数据", + "properties": { + "result": { + "type": "object", + "properties": { + "last_execute_time": {"type": "integer"}, + "task_type": {"type": "string"}, + "task_result": { + "type": "array", + "items": [ + { + "type": "object", + "properties": { + "timed": {"type": "boolean"}, + "rpms": { + "type": "array", + "items": [ + { + "type": "object", + "properties": { + "avaliable_rpm": {"type": "string"}, + "result": {"type": "string"}, + }, + "required": ["avaliable_rpm", "result"], + }, + ], + }, + }, + "required": ["timed", "rpms"], + }, + ], + }, + }, + "required": ["last_execute_time", "task_type", "task_result"], + }, + }, + "required": ["result"], + }, + }, + }, + ), + NodePool( + _id=str(uuid.uuid4()), # 自动生成一个唯一的 ID + service_id="1137ab09-20ae-4278-8346-524d4ce81d2f", + call_id="api", + name="【API】获取漏洞扫描结果", + description="调用接口，获取漏洞扫描结果", + params_schema={ + "full_url": "https://a-ops3.local/vulnerabilities/task/cve_scan/result/get", + "service_id": "aops-apollo", + "method": "post", + "input_data": {}, + "timeout": 300, + "output_key": [{"key": "data", "path": "result"}], + }, + output_schema={ + "type": "object", + "description": "API的返回信息", + "properties": { + "status_code": {"type": "integer", "description": "HTTP状态码"}, + "data": { + "type": "object", + "description": "接口的返回数据", + "properties": { + "result": { + "type": "object", + "properties": { + "last_execute_time": {"type": "integer"}, + "task_type": {"type": "string"}, + "task_result": { + "type": "array", + "items": [ + { + "type": "object", + "properties": { + "timed": {"type": "boolean"}, + "cve_list": { + "type": "array", + "items": [ + { + "type": "object", + "properties": { + "cve_id": {"type": "string"}, + "cve_description": {"type": "string"}, + "rpms": { + "type": "array", + "items": [{"type": "string"}], + }, + "severity": {"type": "string"}, + "cvss_score": {"type": "number"}, + }, + "required": [ + "cve_id", + "cve_description", + "rpms", + "severity", + "cvss_score", + ], + }, + ], + }, + }, + "required": ["timed", "cve_list"], + }, + ], + }, + }, + "required": ["last_execute_time", "task_type", "task_result"], + }, + }, + "required": ["result"], + }, + }, + "required": ["status_code", "data"], + }, + ), + NodePool( + _id=str(uuid.uuid4()), # 自动生成一个唯一的 ID + service_id="1137ab09-20ae-4278-8346-524d4ce81d2f", + call_id="llm", + name="【LLM】生成报告", + description="调用 LLM 生成报告", + params_schema={ + "system_prompt": "你是一个专业的安全专家，擅长生成漏洞相关的分析报告。", + "user_prompt": "请根据历史对话（包括对话中AI助手的回答，和工具的输出），\ +生成一份漏洞相关任务执行的结果报告。\n\n要求如下：\n\ +1. 报告需要包含任务的详细信息，包括任务的名称、创建时间、任务类型等。\n\ +2. 报告需要以“CVE任务执行报告”为标题，且不得超过2000字。\n", + "temperature": 0.7, + "max_tokens": 2048, + }, + output_schema={ + "type": "object", + "properties": {"message": {"type": "string", "description": "大模型的输出内容"}}, + }, + ), ] - collection = MongoDB.get_collection('node') + collection = MongoDB.get_collection("node") result = collection.delete_many({}) # 将 NodePool 模型转换为字典并插入到 MongoDB 中 try: import time + for node_pool in node_pools: time.sleep(1) print(node_pool.service_id) result = collection.update_one( - {'_id': node_pool.id}, # 查找条件，这里假设 _id 即为 node_pool.id + {"_id": node_pool.id}, # 查找条件，这里假设 _id 即为 node_pool.id {"$set": node_pool.dict(by_alias=True)}, # 更新操作 - upsert=True # 如果不存在则插入新文档 + upsert=True, # 如果不存在则插入新文档 ) print("updata success") print(result.upserted_id) - - return result.upserted_id or node_pool.id # 返回插入或更新的文档ID except PyMongoError as e: print(f"An error occurred while inserting the document: {e}") raise # 或者根据需要选择是否重新抛出异常 -class Permission(BaseModel): - """App的权限配置""" - - type: PermissionType = Field( - description="权限类型", default=PermissionType.PRIVATE) - users: list[str] = Field(description="可访问的用户列表", default=[]) - - class PositionItem(BaseModel): """请求/响应中的前端相对位置变量类""" + x: float y: float @@ -330,8 +561,7 @@ class AppFlow(PoolBase): enabled: bool = Field(description="是否启用", default=True) path: str = Field(description="Flow的路径") - focus_point: PositionItem = Field( - description="Flow的视觉焦点", default=PositionItem(x=0, y=0)) + focus_point: PositionItem = Field(description="Flow的视觉焦点", default=PositionItem(x=0, y=0)) class AppLink(BaseModel): @@ -355,14 +585,15 @@ class AppPool(PoolBase): first_questions: list[str] = Field(description="推荐问题", default=[]) history_len: int = Field(3, ge=1, le=10, description="对话轮次（1～10）") permission: Permission = Field( - description="应用权限配置", default=Permission(type=PermissionType.PRIVATE.value, users=[])) + description="应用权限配置", default=Permission(type=PermissionType.PRIVATE, users=[]), + ) flows: list[AppFlow] = Field(description="Flow列表", default=[]) favorites: list[str] = Field(description="收藏此应用的用户列表", default=[]) hashes: dict[str, str] = Field(description="关联文件的hash值", default={}) async def insert_app_pool(): - collection = MongoDB.get_collection('app') + collection = MongoDB.get_collection("app") result = collection.delete_many({}) # 清空集合中的所有文档（仅用于演示） return app_pool = AppPool( @@ -377,28 +608,30 @@ async def insert_app_pool(): first_questions=["What is your name?", "How are you?"], # 提供推荐问题列表 history_len=5, # 设置对话轮次 permission=Permission(type=PermissionType.PUBLIC, users=["user1", "user2"]), # 设置权限配置 - flows=[AppFlow( - _id="test", # 提供一个唯一的标识符 - name="Main Flow", # 提供一个名称 - description="Description of the main flow", # 提供描述 - path="main_flow", - focus_point=PositionItem(x=0.5, y=0.5) - )], # 添加Flows + flows=[ + AppFlow( + _id="test", # 提供一个唯一的标识符 + name="Main Flow", # 提供一个名称 + description="Description of the main flow", # 提供描述 + path="main_flow", + focus_point=PositionItem(x=0.5, y=0.5), + ) + ], # 添加Flows favorites=["user3", "user4"], # 收藏此应用的用户列表 - hashes={"file1": "hash1", "file2": "hash2"} # 关联文件的hash值 + hashes={"file1": "hash1", "file2": "hash2"}, # 关联文件的hash值 ) collection.update_one( - {'_id': app_pool.id}, # 查找条件，这里假设 _id 即为 app_pool.id + {"_id": app_pool.id}, # 查找条件，这里假设 _id 即为 app_pool.id {"$set": app_pool.dict(by_alias=True)}, # 更新操作 - upsert=True # 如果不存在则插入新文档 + upsert=True, # 如果不存在则插入新文档 ) for i in range(100): print(i) app_pool = AppPool( - _id="my_test_"+str(i), # 自动生成一个唯一的 ID + _id="my_test_" + str(i), # 自动生成一个唯一的 ID author="test", # 使用 "author_id" 作为作者ID - name="Example App Pool "+str(i), # 提供名称 - description="This is my test "+str(i), # 提供描述 + name="Example App Pool " + str(i), # 提供名称 + description="This is my test " + str(i), # 提供描述 published=i % 2, icon="icon_url", # 提供图标URL type="example_type", # 提供应用类型 @@ -406,23 +639,25 @@ async def insert_app_pool(): first_questions=["What is your name?", "How are you?"], # 提供推荐问题列表 history_len=5, # 设置对话轮次 permission=Permission(type=PermissionType.PUBLIC, users=["user1", "user2"]), # 设置权限配置 - flows=[AppFlow( - _id="test", # 提供一个唯一的标识符 - name="Main Flow", # 提供一个名称 - description="Description of the main flow", # 提供描述 - path="main_flow", - focus_point=PositionItem(x=0.5, y=0.5) - )], # 添加Flows + flows=[ + AppFlow( + _id="test", # 提供一个唯一的标识符 + name="Main Flow", # 提供一个名称 + description="Description of the main flow", # 提供描述 + path="main_flow", + focus_point=PositionItem(x=0.5, y=0.5), + ) + ], # 添加Flows favorites=["user3", "user4"], # 收藏此应用的用户列表 - hashes={"file1": "hash1", "file2": "hash2"} # 关联文件的hash值 + hashes={"file1": "hash1", "file2": "hash2"}, # 关联文件的hash值 ) # 将 AppPool 模型转换为字典并插入到 MongoDB 中 try: result = collection.update_one( - {'_id': app_pool.id}, # 查找条件，这里假设 _id 即为 app_pool.id + {"_id": app_pool.id}, # 查找条件，这里假设 _id 即为 app_pool.id {"$set": app_pool.dict(by_alias=True)}, # 更新操作 - upsert=True # 如果不存在则插入新文档 + upsert=True, # 如果不存在则插入新文档 ) print(f"Inserted or updated document with id: {app_pool.id}") except PyMongoError as e: @@ -431,10 +666,10 @@ async def insert_app_pool(): for i in range(100): print(i) app_pool = AppPool( - _id="test_"+str(i), # 自动生成一个唯一的 ID - author="test_"+str(i), # 使用 "author_id" 作为作者ID + _id="test_" + str(i), # 自动生成一个唯一的 ID + author="test_" + str(i), # 使用 "author_id" 作为作者ID name="Example App Pool", # 提供名称 - description="This is test "+str(i), # 提供描述 + description="This is test " + str(i), # 提供描述 published=i % 2, icon="icon_url", # 提供图标URL type="example_type", # 提供应用类型 @@ -442,23 +677,25 @@ async def insert_app_pool(): first_questions=["What is your name?", "How are you?"], # 提供推荐问题列表 history_len=5, # 设置对话轮次 permission=Permission(type=PermissionType.PUBLIC, users=["user1", "user2"]), # 设置权限配置 - flows=[AppFlow( - _id="test", # 提供一个唯一的标识符 - name="Main Flow", # 提供一个名称 - description="Description of the main flow", # 提供描述 - path="main_flow", - focus_point=PositionItem(x=0.5, y=0.5) - )], # 添加Flows + flows=[ + AppFlow( + _id="test", # 提供一个唯一的标识符 + name="Main Flow", # 提供一个名称 + description="Description of the main flow", # 提供描述 + path="main_flow", + focus_point=PositionItem(x=0.5, y=0.5), + ) + ], # 添加Flows favorites=["user3", "user4"], # 收藏此应用的用户列表 - hashes={"file1": "hash1", "file2": "hash2"} # 关联文件的hash值 + hashes={"file1": "hash1", "file2": "hash2"}, # 关联文件的hash值 ) # 将 AppPool 模型转换为字典并插入到 MongoDB 中 try: result = collection.update_one( - {'_id': app_pool.id}, # 查找条件，这里假设 _id 即为 app_pool.id + {"_id": app_pool.id}, # 查找条件，这里假设 _id 即为 app_pool.id {"$set": app_pool.dict(by_alias=True)}, # 更新操作 - upsert=True # 如果不存在则插入新文档 + upsert=True, # 如果不存在则插入新文档 ) print(f"Inserted or updated document with id: {app_pool.id}") except PyMongoError as e: @@ -466,7 +703,7 @@ async def insert_app_pool(): raise # 或者根据需要选择是否重新抛出异常 -def query_all_target(tag:str): +def query_all_target(tag: str): """查询所有插入到MongoDB中的node数据""" collection = MongoDB.get_collection(tag) try: @@ -476,8 +713,9 @@ def query_all_target(tag:str): except PyMongoError as e: print(f"An error occurred while querying the documents: {e}") + # 使用 asyncio 运行异步函数 if __name__ == "__main__": - # asyncio.run(insert_node_pool()) asyncio.run(insert_service_pool()) - query_all_target("node") \ No newline at end of file + asyncio.run(insert_node_pool()) + query_all_target("node") -- Gitee From 31a84a1f338f1f0423038406452f718d5c3b0688 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E5=8F=B2=E9=B8=BF=E5=AE=87?= Date: Tue, 18 Feb 2025 15:29:50 +0800 Subject: [PATCH 2/2] =?UTF-8?q?mock:=20=E6=9B=B4=E6=96=B0=20A-Ops=20?= =?UTF-8?q?=E6=8E=A5=E5=8F=A3=20ID?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: 史鸿宇 --- mock/make_data.py | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/mock/make_data.py b/mock/make_data.py index 7d8ee314a..5b8aa0757 100644 --- a/mock/make_data.py +++ b/mock/make_data.py @@ -276,7 +276,7 @@ async def insert_node_pool() -> None: }, ), NodePool( - _id=str(uuid.uuid4()), # 自动生成一个唯一的 ID + _id="343da7db-5da8-42ef-9b59-cc56df54d9aa", service_id="1137ab09-20ae-4278-8346-524d4ce81d2f", call_id="api", name="【API】获取任务简介", @@ -339,7 +339,7 @@ async def insert_node_pool() -> None: }, ), NodePool( - _id=str(uuid.uuid4()), # 自动生成一个唯一的 ID + _id="8841e328-da5b-45c7-8839-5b8054a92de7", service_id="1137ab09-20ae-4278-8346-524d4ce81d2f", call_id="choice", name="【CHOICE】判断任务类型", @@ -366,7 +366,7 @@ async def insert_node_pool() -> None: }, ), NodePool( - _id=str(uuid.uuid4()), # 自动生成一个唯一的 ID + _id="7377ad0d-f867-46fe-806a-d0c4535d2f1a", service_id="1137ab09-20ae-4278-8346-524d4ce81d2f", call_id="api", name="【API】获取CVE修复的结果", @@ -433,7 +433,7 @@ async def insert_node_pool() -> None: }, ), NodePool( - _id=str(uuid.uuid4()), # 自动生成一个唯一的 ID + _id="3d94b288-a0df-4717-b75c-fc2c67e24294", service_id="1137ab09-20ae-4278-8346-524d4ce81d2f", call_id="api", name="【API】获取漏洞扫描结果", @@ -508,7 +508,7 @@ async def insert_node_pool() -> None: }, ), NodePool( - _id=str(uuid.uuid4()), # 自动生成一个唯一的 ID + _id="1a8ddfb9-c894-4819-ab9b-88fcb5f14c10", service_id="1137ab09-20ae-4278-8346-524d4ce81d2f", call_id="llm", name="【LLM】生成报告", -- Gitee