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()

muagent/sandbox/nbclient.py

@@ -24,9 +24,9 @@
 
 
 
-
-import logging
-logging.basicConfig(level=logging.INFO)
+from ..utils.logger import logging
+# import logging
+# logging.basicConfig(level=logging.INFO)
 
 
 

muagent/schemas/kb/__init__.py

@@ -1,13 +1,14 @@
 # lingsi 游走推理服务 old
 import json
 import requests
-import logging
+# import logging
+from ..utils.logger import logging
 import copy
 import sys
 import os
 
-logger = logging.getLogger()
-logging.basicConfig(level=logging.DEBUG)
+# logger = logging.getLogger()
+# logging.basicConfig(level=logging.DEBUG)
 
 #step 1  retriver
 

muagent/service/ekg_reasoning/__init__.py

@@ -28,7 +28,8 @@
 import json
 import requests
 import time
-import logging
+# import logging
+from ..utils.logger import logging
 import copy
 import sys
 import os,  base64

muagent/service/ekg_reasoning/src/__init__.py

@@ -24,7 +24,8 @@
 import json
 import requests
 import time
-import logging
+# import logging
+from ..utils.logger import logging
 import copy
 import sys
 import os,  base64
@@ -66,8 +67,8 @@
 
 
 
-# 配置logging模块
-logging.basicConfig(format='%(asctime)s - %(levelname)s - %(message)s - %(lineno)d', level=logging.INFO)
+# # 配置logging模块
+# logging.basicConfig(format='%(asctime)s - %(levelname)s - %(message)s - %(lineno)d', level=logging.INFO)
 
 
 

muagent/service/ekg_reasoning/src/geabase_handler/__init__.py

@@ -29,7 +29,8 @@
 import json
 import requests
 import time
-import logging
+# import logging
+from ..utils.logger import logging
 import copy
 import sys
 import os,  base64

muagent/service/ekg_reasoning/src/geabase_handler/geabase_handlerplus.py

@@ -26,7 +26,8 @@
 import json
 import requests
 import time
-import logging
+# import logging
+from ..utils.logger import logging
 import copy
 import sys
 import os,  base64
@@ -63,8 +64,8 @@
 
 
 
-# 配置logging模块
-logging.basicConfig(format='%(asctime)s - %(levelname)s - %(message)s - %(lineno)d', level=logging.INFO)
+# # 配置logging模块
+# logging.basicConfig(format='%(asctime)s - %(levelname)s - %(message)s - %(lineno)d', level=logging.INFO)
 
 
 

muagent/service/ekg_reasoning/src/graph_search/__init__.py

@@ -4,9 +4,9 @@
 import logging
 import copy
 import os
-
-logger = logging.getLogger()
-logging.basicConfig(level=logging.DEBUG)
+from ..utils.logger import logging
+# logger = logging.getLogger()
+# logging.basicConfig(level=logging.DEBUG)
 RETRY_MAX_NUM = 3
 def intention_recognition_ekgfunc( root_node_id, rule, query, memory, start_from_root = True,
         url= os.environ['intention_url'] ):

muagent/service/ekg_reasoning/src/graph_search/call_old_fuction.py

@@ -22,8 +22,9 @@
 
 from muagent.connector.schema import Message
 from src.utils.call_llm import call_llm,extract_final_result
-import logging
-logging.basicConfig(level=logging.INFO)
+from ..utils.logger import logging
+# import logging
+# logging.basicConfig(level=logging.INFO)
 
 # from loguru import logger as logging
 from src.geabase_handler.geabase_handlerplus import GB_handler

muagent/service/ekg_reasoning/src/graph_search/geabase_search_plus.py

@@ -17,8 +17,9 @@
 from muagent.connector.schema import Message
 from src.utils.call_llm import call_llm, extract_final_result, robust_call_llm
 
-import logging
-logging.basicConfig(level=logging.INFO)
+from ..utils.logger import logging
+# import logging
+# logging.basicConfig(level=logging.INFO)
 
 from src.geabase_handler.geabase_handlerplus import GB_handler
 from src.utils.normalize import hash_id

muagent/service/ekg_reasoning/src/graph_search/graph_search_main.py

@@ -21,8 +21,8 @@
 from colorama import Fore
 from Crypto.Cipher import AES
 from loguru import logger
-import logging
-
+# import logging
+from loguru import logger as logging
 
 
 MOST_RETRY_TIMES = 5

muagent/service/ekg_reasoning/src/graph_search/task_node_agent.py

@@ -10,6 +10,14 @@ def get_logger():
     logger_loggering = logging.getLogger()
     return logger_loggering
 
+import os
+if os.environ.get("operation_mode")=="antcode":
+    import logging
+    FORMAT = '[%(asctime)s %(filename)s->%(funcName)s():%(lineno)s]%(levelname)s: %(message)s'
+    logging.basicConfig(level=logging.DEBUG, format=FORMAT)
+else:
+    from loguru import logger as logging
+
 
 if __name__ == '__main__':
     logger = get_logger()

muagent/service/ekg_reasoning/src/graphstructure/__init__.py

@@ -5,10 +5,10 @@
 
 setuptools.setup(
     name="codefuse-muagent",
-    version="0.1.0",
+    version="0.1.1",
     author="shanshi",
     author_email="wyp311395@antgroup.com",
-    description="A multi-agent framework that facilitates the rapid construction of collaborative teams of agents.",
+    description="An Innovative Agent Framework Driven by KG Engine",
     # long_description=long_description,
     long_description_content_type="text/markdown",
     url="https://github.com/codefuse-ai/CodeFuse-muAgent",
@@ -40,7 +40,7 @@
         "ollama",
         "colorama",
         "pycryptodome",
-        "dashscope"
+        "dashscope",
         # 
         "chromadb==0.4.17",
         "javalang==0.13.0",

muagent/service/ekg_reasoning/src/graphstructure/graphstrcturesearchfun.py

@@ -19,9 +19,10 @@
 )
 sys.path.append(src_dir)
 
-import logging
+# import logging
+from loguru import logger as logging
 # Set the logging level to WARNING, which will suppress INFO and DEBUG messages
-logging.basicConfig(level=logging.ERROR)
+# logging.basicConfig(level=logging.ERROR)
 
 
 from muagent import EKG, get_ekg_project_config_from_env

muagent/service/ekg_reasoning/src/intention_recognition/__init__.py

@@ -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'
 
 ########################################
 ########## 以下参数暂不涉及无需配置 ########