Merge pull request #101 from codefuse-ai/muagent_dev

[muagent v0.1.1] add reasoning logger and shutdown nbclient kernel

Author: lightislost

examples/test_config.py.example

@@ -1,5 +1,6 @@
 import os, openai, base64
 from loguru import logger
+import json
 
 os.environ["DM_llm_name"] = 'Qwen2_72B_Instruct_OpsGPT'  #or  gpt_4
 
@@ -123,7 +124,7 @@ DB_CONFIGS = {
     "gb_config": {
         "gb_type": "NebulaHandler",
         "extra_kwargs": {
-            'host':'graphd',
+            'host': os.environ['nb_host'],
             'port': '9669',
             'username': os.environ['nb_username'],
             'password': os.environ['nb_password'],
@@ -133,7 +134,7 @@ DB_CONFIGS = {
     "tb_config": {
         "tb_type": 'TBaseHandler',
         "index_name": "opsgptkg",
-        "host": 'redis-stack',
+        "host": os.environ['tb_host'],
         "port": '6379',
         "username": os.environ['tb_username'],
         "password": os.environ['tb_password'],
@@ -146,6 +147,7 @@ DB_CONFIGS = {
 os.environ["DB_CONFIGS"] = json.dumps(DB_CONFIGS)
 
 
+os.environ["clear_history_data"] = "False" # 'True'
 
 ########################################
 ########## 以下参数暂不涉及无需配置 ########

muagent/db_handler/graph_db_handler/geabase_handler.py

@@ -7,7 +7,8 @@
     from gdbc2.geabase_client import GeaBaseClient, Node, Edge, MutateBatchOperation, GeaBaseUtil
     from gdbc2.geabase_env import GeaBaseEnv
 except:
-    logger.error("ignore this sdk")
+    pass
+    # logger.error("ignore this sdk")
 
 from .base_gb_handler import GBHandler
 from muagent.db_handler.utils import deduplicate_dict

muagent/db_handler/graph_db_handler/nebula_handler.py

@@ -1021,7 +1021,7 @@ def is_nodeid_exist(self, nodeid: str) -> bool:
     def check_edge_exist(self, src_id: str, dst_id: str) -> bool:
         try:
             if self.get_current_edge(src_id, dst_id):
-                print('Edge already exists!')
+                # print('Edge already exists!')
                 return True
         except Exception as e:
             print('Edge not exists!')
@@ -1041,7 +1041,7 @@ def check_node_before_execute(self, nodeid: str, action: str = ""):
             )
         # 节点是否存在: 1.删除时节点存在不报错 2.添加时节点存在报错 3.update时节点存在不报错
         if self.is_nodeid_exist(nodeid) and action == "add":
-            logger.info('Node already exits!')
+            # logger.info('Node already exits!')
             return GbaseExecStatus(
             errorMessage='GDB_ENGINE_PRIMARY_KEY_DUPLICATE', 
             errorCode=1, 

muagent/ekg_project.py

@@ -15,6 +15,9 @@
 import time
 import random
 
+os.environ["operation_mode"] = "open_source"
+os.environ["intention_url"] = ""
+
 from .llm_models import LLMConfig, EmbedConfig
 from .schemas.db import TBConfig, GBConfig
 from .schemas.models import ModelConfig
@@ -228,7 +231,6 @@ def get_ekg_project_config_from_env(
         model_type = random.choice(embedding_list)
         default_model_config = project_configs["model_configs"][model_type]
         project_configs["model_configs"]["default_embed"] = default_model_config
-    project_configs[k] = v
 
     # init embedding configs
     if embed_configs:
@@ -286,7 +288,7 @@ def get_ekg_project_config_from_env(
         project_configs["prompt_configs"] = prompt_configs
     else:
         logger.warning(
-            f"Cant't init any AGENT_CONFIGS in this env."
+            f"Cant't init any PROMPT_CONFIGS in this env."
         ) 
 
 

muagent/prompt_manager/language/__init__.py

@@ -179,7 +179,8 @@ def reset_notebook(self) -> str:
         asyncio.run(self._restart_client())
         return "Reset notebook"
 
-
+    async def close_client(self):
+        await self.nb_client.km.shutdown_kernel(now=True)
 
 
 
@@ -291,7 +292,7 @@ def restart(self, ) -> CodeBoxStatus:
         return CodeBoxStatus(status="restared")
 
     def stop(self, ) -> CodeBoxStatus:
-        pass
+        asyncio.run(self.nbe.close_client())
 
     def __del__(self):
         self.stop()