agentUniverse 0.0.17.dev5__tar.gz → 0.0.17.dev6__tar.gz

{agentuniverse-0.0.17.dev5 → agentuniverse-0.0.17.dev6}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: agentUniverse
-Version: 0.0.17.dev5
+Version: 0.0.17.dev6
 Summary: agentUniverse is a framework for developing applications powered by multi-agent base on large language model.
 Author: AntGroup
 Author-email: jerry.zzw@antgroup.com
@@ -34,7 +34,7 @@ Requires-Dist: langchain-anthropic (>=0.1.13,<0.2.0)
 Requires-Dist: langchain-community (==0.0.38)
 Requires-Dist: langchain-core (==0.1.52)
 Requires-Dist: loguru (==0.7.2)
-Requires-Dist: mcp (>=1.8.0,<2.0.0)
+Requires-Dist: mcp (>=1.9.0,<1.10.0)
 Requires-Dist: myst-parser (>=2.0.0,<3.0.0)
 Requires-Dist: networkx (>=3.3,<4.0)
 Requires-Dist: numpy (>=1.26.0,<2.0.0)

{agentuniverse-0.0.17.dev5 → agentuniverse-0.0.17.dev6}/agentuniverse/agent/agent.py

@@ -44,10 +44,9 @@ from agentuniverse.base.context.framework_context_manager import \
 from agentuniverse.base.util.agent_util import process_agent_llm_config
 from agentuniverse.base.util.common_util import stream_output
 from agentuniverse.base.util.logging.logging_util import LOGGER
-from agentuniverse.base.util.memory_util import generate_messages, \
-    get_memory_string
-from agentuniverse.base.util.system_util import process_dict_with_funcs, \
-    is_system_builtin
+from agentuniverse.base.util.memory_util import generate_messages, get_memory_string
+from agentuniverse.base.util.system_util import process_dict_with_funcs, is_system_builtin
+from agentuniverse.base.util.tracing.au_trace_manager import AuTraceManager
 from agentuniverse.llm.llm import LLM
 from agentuniverse.llm.llm_manager import LLMManager
 from agentuniverse.prompt.chat_prompt import ChatPrompt
@@ -98,6 +97,15 @@ class Agent(ComponentBase, ABC):
         """
         pass
 
+    def update_trace_context(self, input_object: InputObject):
+        session_id = input_object.get_data("session_id")
+        if session_id:
+            AuTraceManager().trace_context.set_session_id(session_id)
+
+        trace_id = input_object.get_data("trace_id")
+        if trace_id:
+            AuTraceManager().trace_context.set_trace_id(trace_id)
+
     @trace_agent
     def run(self, **kwargs) -> OutputObject:
         """Agent instance running entry.
@@ -108,6 +116,8 @@ class Agent(ComponentBase, ABC):
         self.input_check(kwargs)
         input_object = InputObject(kwargs)
 
+        self.update_trace_context(input_object)
+
         agent_input = self.pre_parse_input(input_object)
 
         planner_result = self.execute(input_object, agent_input)
@@ -122,6 +132,7 @@ class Agent(ComponentBase, ABC):
     async def async_run(self, **kwargs) -> OutputObject:
         self.input_check(kwargs)
         input_object = InputObject(kwargs)
+        self.update_trace_context(input_object)
 
         agent_input = self.pre_parse_input(input_object)
 
@@ -389,7 +400,7 @@ class Agent(ComponentBase, ABC):
                 tool_input = {key: input_object.get_data(key) for key in tool.input_keys}
                 tool_results.append(str(tool.run(**tool_input)))
             except:
-                pass
+                LOGGER.warn(f'Tool {tool_name} call failed, maybe invalid or lack arguments')
         return "\n\n".join(tool_results)
 
     async def async_invoke_tools(self, input_object: InputObject, **kwargs) -> str:

{agentuniverse-0.0.17.dev5 → agentuniverse-0.0.17.dev6}/agentuniverse/agent_serve/web/mcp/mcp_server_manager.py

@@ -114,7 +114,7 @@ class MCPServerManager:
         for server_name, tool_dict in self.server_tool_map.items():
             mcp_server = FastMCP(server_name)
 
-            for _tool in tool_dict['tool']:
+            for _tool in set(tool_dict['tool']):
                 tool_instance = ToolManager().get_instance_obj(_tool)
                 mcp_server.add_tool(
                     fn=create_exposed_tool_method_wrapper(_tool),
@@ -122,7 +122,7 @@ class MCPServerManager:
                     description=tool_instance.description
                 )
 
-            for _toolkit in tool_dict['toolkit']:
+            for _toolkit in set(tool_dict['toolkit']):
                 toolkit_instance = ToolkitManager().get_instance_obj(_toolkit)
                 for _tool in toolkit_instance.tool_names:
                     tool_instance = ToolManager().get_instance_obj(_tool)
@@ -136,7 +136,7 @@ class MCPServerManager:
                 'server': mcp_server
             }
             if transport == 'sse':
-                mcp_server_dict['mount'] = Mount(f"/{server_name}", app=mcp_server.sse_app(f"/{server_name}"))
+                mcp_server_dict['mount'] = Mount(f"/{server_name}", app=mcp_server.sse_app())
 
             mcp_server_list.append(mcp_server_dict)
 

{agentuniverse-0.0.17.dev5 → agentuniverse-0.0.17.dev6}/agentuniverse/base/annotation/trace.py

@@ -14,6 +14,7 @@ import uuid
 from functools import wraps
 
 from agentuniverse.agent.memory.conversation_memory.conversation_memory_module import ConversationMemoryModule
+from agentuniverse.base.config.application_configer.application_config_manager import ApplicationConfigManager
 from agentuniverse.base.util.monitor.monitor import Monitor
 from agentuniverse.llm.llm_output import LLMOutput
 
@@ -23,6 +24,12 @@ def trace_llm(func):
 
     Decorator to trace the LLM invocation, add llm input and output to the monitor.
     """
+    def plugins(func):
+        llm_plugins = ApplicationConfigManager().app_configer.llm_plugins
+        warp_func = func
+        for item in llm_plugins:
+            warp_func = item(func)
+        return warp_func
 
     @wraps(func)
     async def wrapper_async(*args, **kwargs):
@@ -41,7 +48,7 @@ def trace_llm(func):
 
         if self and hasattr(self, 'tracing'):
             if self.tracing is False:
-                return await func(*args, **kwargs)
+                return await plugins(func)(*args, **kwargs)
 
         # add invocation chain to the monitor module.
         Monitor.add_invocation_chain({'source': source, 'type': 'llm'})
@@ -50,7 +57,7 @@ def trace_llm(func):
         Monitor().trace_llm_input(source=source, llm_input=llm_input)
 
         # invoke function
-        result = await func(*args, **kwargs)
+        result = await plugins(func)(*args, **kwargs)
         # not streaming
         if isinstance(result, LLMOutput):
             # add llm invocation info to monitor
@@ -95,16 +102,15 @@ def trace_llm(func):
 
         if self and hasattr(self, 'tracing'):
             if self.tracing is False:
-                return func(*args, **kwargs)
+                return plugins(func)(*args, **kwargs)
 
         # add invocation chain to the monitor module.
         Monitor.add_invocation_chain({'source': source, 'type': 'llm'})
 
         start_time = time.time()
         Monitor().trace_llm_input(source=source, llm_input=llm_input)
-
         # invoke function
-        result = func(*args, **kwargs)
+        result = plugins(func)(*args, **kwargs)
         # not streaming
         if isinstance(result, LLMOutput):
             # add llm invocation info to monitor
@@ -134,10 +140,9 @@ def trace_llm(func):
             return gen_iterator()
 
     if asyncio.iscoroutinefunction(func):
-        # async function
+
         return wrapper_async
     else:
-        # sync function
         return wrapper_sync
 
 
@@ -182,16 +187,16 @@ def trace_agent(func):
         start_info = get_caller_info()
         pair_id = f"agent_{uuid.uuid4().hex}"
         kwargs['memory_source_info'] = start_info
+        # add invocation chain to the monitor module.
+        Monitor.init_invocation_chain()
+        Monitor.add_invocation_chain({'source': source, 'type': 'agent'})
         ConversationMemoryModule().add_agent_input_info(start_info, self, agent_input, pair_id)
         if tracing is False:
             result = await func(*args, **kwargs)
             ConversationMemoryModule().add_agent_result_info(self, result, start_info, pair_id)
+            Monitor.pop_invocation_chain()
             return result
 
-        # add invocation chain to the monitor module.
-        Monitor.init_invocation_chain()
-        Monitor.add_invocation_chain({'source': source, 'type': 'agent'})
-
         start_time = time.time()
         Monitor().trace_agent_input(source=source, agent_input=agent_input)
 
@@ -225,16 +230,16 @@ def trace_agent(func):
         pair_id = f"agent_{uuid.uuid4().hex}"
         start_info = get_caller_info()
         kwargs['memory_source_info'] = start_info
+        # add invocation chain to the monitor module.
+        Monitor.init_invocation_chain()
+        Monitor.add_invocation_chain({'source': source, 'type': 'agent'})
         ConversationMemoryModule().add_agent_input_info(start_info, self, agent_input, pair_id)
         if tracing is False:
             result = func(*args, **kwargs)
             ConversationMemoryModule().add_agent_result_info(self, result, start_info, pair_id)
+            Monitor.pop_invocation_chain()
             return result
 
-        # add invocation chain to the monitor module.
-        Monitor.init_invocation_chain()
-        Monitor.add_invocation_chain({'source': source, 'type': 'agent'})
-
         start_time = time.time()
         Monitor().trace_agent_input(source=source, agent_input=agent_input)
 

{agentuniverse-0.0.17.dev5 → agentuniverse-0.0.17.dev6}/agentuniverse/base/config/application_configer/app_configer.py

@@ -1,12 +1,11 @@
 # !/usr/bin/env python3
 # -*- coding:utf-8 -*-
-
 # @Time    : 2024/3/12 16:17
 # @Author  : jerry.zzw 
 # @Email   : jerry.zzw@antgroup.com
 # @FileName: app_configer.py
-from typing import Optional, Dict
-
+import importlib
+from typing import Optional, Dict, Any
 from agentuniverse.base.config.component_configer.component_configer import \
     ComponentConfiger
 from agentuniverse.base.config.component_configer.configers.llm_configer import LLMConfiger
@@ -56,6 +55,7 @@ class AppConfiger(object):
         self.__agent_llm_set: Optional[set[str]] = set()
         self.__agent_tool_set: Optional[set[str]] = set()
         self.__agent_toolkit_set: Optional[set[str]] = set()
+        self.__llm_plugins: Optional[Any] = set()
 
     @property
     def base_info_appname(self) -> Optional[str]:
@@ -247,6 +247,19 @@ class AppConfiger(object):
     def agent_tool_set(self, value: set):
         self.__agent_tool_set = value
 
+    @property
+    def llm_plugins(self):
+        return self.__llm_plugins
+
+    @classmethod
+    def load_llm_plugins(cls, plugin_modules):
+        funcs = []
+        for item in plugin_modules:
+            module_name, func_name = item.rsplit('.', 1)
+            module = importlib.import_module(module_name)
+            funcs.append(getattr(module, func_name))
+        return funcs
+
     def load_by_configer(self, configer: Configer) -> 'AppConfiger':
         """Load the AppConfiger by the given Configer.
 
@@ -283,4 +296,5 @@ class AppConfiger(object):
         self.__core_log_sink_package_list = configer.value.get('CORE_PACKAGE', {}).get('log_sink')
         self.__core_llm_channel_package_list = configer.value.get('CORE_PACKAGE', {}).get('llm_channel')
         self.__conversation_memory_configer = configer.value.get('CONVERSATION_MEMORY', {})
+        self.__llm_plugins = self.load_llm_plugins(configer.value.get("PLUGINS", {}).get("llm_plugins", []))
         return self

{agentuniverse-0.0.17.dev5 → agentuniverse-0.0.17.dev6}/agentuniverse/llm/llm.py

@@ -8,6 +8,8 @@
 from abc import abstractmethod
 from copy import deepcopy
 from typing import Optional, Any, AsyncIterator, Iterator, Union
+
+import tiktoken
 from langchain_core.language_models.base import BaseLanguageModel
 from langchain_core.runnables import Runnable
 
@@ -177,10 +179,17 @@ class LLM(ComponentBase):
 
         Args:
             text: The string input to tokenize.
+            model: The model you want to calculate
+
 
         Returns:
             The integer number of tokens in the text.
         """
+        try:
+            encoding = tiktoken.encoding_for_model(self.model_name)
+        except KeyError:
+            encoding = tiktoken.get_encoding("cl100k_base")
+        return len(encoding.encode(text))
 
     def as_langchain_runnable(self, params=None) -> Runnable:
         """Get the langchain llm class."""
@@ -194,7 +203,7 @@ class LLM(ComponentBase):
         try:
             self.init_channel()
             if self._channel_instance:
-                return self._channel_instance.call(*args, **kwargs)
+                return self._channel_instance._call(*args, **kwargs)
             return self._call(*args, **kwargs)
         except Exception as e:
             LOGGER.error(f'Error in LLM call: {e}')
@@ -206,7 +215,7 @@ class LLM(ComponentBase):
         try:
             self.init_channel()
             if self._channel_instance:
-                return await self._channel_instance.acall(*args, **kwargs)
+                return await self._channel_instance._acall(*args, **kwargs)
             return await self._acall(*args, **kwargs)
         except Exception as e:
             LOGGER.error(f'Error in LLM acall: {e}')

{agentuniverse-0.0.17.dev5 → agentuniverse-0.0.17.dev6}/agentuniverse/llm/llm_channel/llm_channel.py

@@ -32,6 +32,8 @@ class LLMChannel(ComponentBase):
     channel_proxy: Optional[str] = None
     channel_model_name: Optional[str] = None
     channel_ext_info: Optional[dict] = None
+    ext_headers: Optional[dict] = {}
+    ext_params: Optional[dict] = {}
 
     model_support_stream: Optional[bool] = None
     model_support_max_context_length: Optional[int] = None
@@ -68,6 +70,19 @@ class LLMChannel(ComponentBase):
             self.model_support_max_tokens = component_configer.model_support_max_tokens
         if hasattr(component_configer, "model_is_openai_protocol_compatible"):
             self.model_is_openai_protocol_compatible = component_configer.model_is_openai_protocol_compatible
+        if component_configer.configer.value.get("ext_headers"):
+            self.ext_headers = component_configer.configer.value.get("extra_headers", {})
+        if component_configer.configer.value.get("ext_params"):
+            self.ext_params = component_configer.configer.value.get("extra_params", {})
+            self.ext_params["stream_options"] = {
+                "include_usage": True
+            }
+        else:
+            self.ext_params = {
+                "stream_options": {
+                    "include_usage": True
+                }
+            }
         return self
 
     def create_copy(self):
@@ -113,19 +128,24 @@ class LLMChannel(ComponentBase):
             streaming = kwargs.pop('stream')
         if self.model_support_stream is False and streaming is True:
             streaming = False
-
         support_max_tokens = self.model_support_max_tokens
-        max_tokens = kwargs.pop('max_tokens', None) or self.channel_model_config.get('max_tokens', None) or support_max_tokens
+        max_tokens = kwargs.pop('max_tokens', None) or self.channel_model_config.get('max_tokens',
+                                                                                     None) or support_max_tokens
         if support_max_tokens:
             max_tokens = min(support_max_tokens, max_tokens)
 
-        self.client = self._new_client()
+        ext_params = self.ext_params.copy()
+        if not streaming:
+            ext_params.pop("stream_options", "")
+        self.client = self._new_client(kwargs.pop("api_base", None))
         chat_completion = self.client.chat.completions.create(
             messages=messages,
             model=kwargs.pop('model', self.channel_model_name),
             temperature=kwargs.pop('temperature', self.channel_model_config.get('temperature')),
             stream=kwargs.pop('stream', streaming),
             max_tokens=max_tokens,
+            extra_body=ext_params,
+            extra_headers=kwargs.pop("extra_headers", self.ext_headers),
             **kwargs,
         )
         if not streaming:
@@ -143,17 +163,22 @@ class LLMChannel(ComponentBase):
             streaming = False
 
         support_max_tokens = self.model_support_max_tokens
-        max_tokens = kwargs.pop('max_tokens', None) or self.channel_model_config.get('max_tokens', None) or support_max_tokens
+        max_tokens = kwargs.pop('max_tokens', None) or self.channel_model_config.get('max_tokens',
+                                                                                     None) or support_max_tokens
         if support_max_tokens:
             max_tokens = min(support_max_tokens, max_tokens)
-
-        self.async_client = self._new_async_client()
+        ext_params = self.ext_params.copy()
+        if not streaming:
+            ext_params.pop("stream_options")
+        self.async_client = self._new_async_client(kwargs.pop("api_base", None))
         chat_completion = await self.async_client.chat.completions.create(
             messages=messages,
             model=kwargs.pop('model', self.channel_model_name),
             temperature=kwargs.pop('temperature', self.channel_model_config.get('temperature')),
             stream=kwargs.pop('stream', streaming),
             max_tokens=max_tokens,
+            extra_headers=kwargs.pop("extra_headers", self.ext_headers),
+            extra_body=ext_params,
             **kwargs,
         )
         if not streaming:
@@ -167,7 +192,7 @@ class LLMChannel(ComponentBase):
         """Convert to the langchain llm class."""
         return DefaultChannelLangchainInstance(self)
 
-    def get_num_tokens(self, text: str) -> int:
+    def get_num_tokens(self, text: str, model=None) -> int:
         """Get the number of tokens present in the text.
 
         Useful for checking if an input will fit in an openai model's context window.
@@ -185,28 +210,28 @@ class LLMChannel(ComponentBase):
     def max_context_length(self) -> int:
         return self.channel_model_config.get('max_context_length')
 
-    def _new_client(self):
+    def _new_client(self, api_base: str = None):
         """Initialize the openai client."""
         if self.client is not None:
             return self.client
         return OpenAI(
             api_key=self.channel_api_key,
             organization=self.channel_organization,
-            base_url=self.channel_api_base,
+            base_url=api_base if api_base else self.channel_api_base,
             timeout=self.channel_model_config.get('request_timeout'),
             max_retries=self.channel_model_config.get('max_retries'),
             http_client=httpx.Client(proxy=self.channel_proxy) if self.channel_proxy else None,
             **(self.channel_model_config.get('client_args') or {}),
         )
 
-    def _new_async_client(self):
+    def _new_async_client(self, api_base=None):
         """Initialize the openai async client."""
         if self.async_client is not None:
             return self.async_client
         return AsyncOpenAI(
             api_key=self.channel_api_key,
             organization=self.channel_organization,
-            base_url=self.channel_api_base,
+            base_url=api_base if api_base else self.channel_api_base,
             timeout=self.channel_model_config.get('request_timeout'),
             max_retries=self.channel_model_config.get('max_retries'),
             http_client=httpx.AsyncClient(proxy=self.channel_proxy) if self.channel_proxy else None,
@@ -234,7 +259,7 @@ class LLMChannel(ComponentBase):
         if not isinstance(chunk, dict):
             chunk = chunk.dict()
         if len(chunk["choices"]) == 0:
-            return
+            return LLMOutput(text="", raw=chat_completion.model_dump())
         choice = chunk["choices"][0]
         message = choice.get("delta")
         text = message.get("content")

{agentuniverse-0.0.17.dev5 → agentuniverse-0.0.17.dev6}/agentuniverse/llm/openai_style_llm.py

@@ -41,29 +41,32 @@ class OpenAIStyleLLM(LLM):
     api_base: Optional[str] = None
     proxy: Optional[str] = None
     client_args: Optional[dict] = None
+    ext_params: Optional[dict] = {}
+    ext_headers: Optional[dict] = {}
 
-    def _new_client(self):
+    def _new_client(self, api_base=None):
         """Initialize the openai client."""
         if self.client is not None:
             return self.client
         return OpenAI(
             api_key=self.api_key,
             organization=self.organization,
-            base_url=self.api_base,
+            base_url=api_base if api_base else self.api_base,
             timeout=self.request_timeout,
             max_retries=self.max_retries,
             http_client=httpx.Client(proxy=self.proxy) if self.proxy else None,
+
             **(self.client_args or {}),
         )
 
-    def _new_async_client(self):
+    def _new_async_client(self, api_base: str = None):
         """Initialize the openai async client."""
         if self.async_client is not None:
             return self.async_client
         return AsyncOpenAI(
             api_key=self.api_key,
             organization=self.organization,
-            base_url=self.api_base,
+            base_url=api_base if api_base else self.api_base,
             timeout=self.request_timeout,
             max_retries=self.max_retries,
             http_client=httpx.AsyncClient(proxy=self.proxy) if self.proxy else None,
@@ -80,7 +83,12 @@ class OpenAIStyleLLM(LLM):
         streaming = kwargs.pop("streaming") if "streaming" in kwargs else self.streaming
         if 'stream' in kwargs:
             streaming = kwargs.pop('stream')
-        self.client = self._new_client()
+        self.client = self._new_client(kwargs.pop("api_base", None))
+        ext_params = self.ext_params.copy()
+        if streaming and "stream_options" not in ext_params:
+            ext_params["stream_options"] = {
+                "include_usage": True
+            }
         client = self.client
         chat_completion = client.chat.completions.create(
             messages=messages,
@@ -88,6 +96,8 @@ class OpenAIStyleLLM(LLM):
             temperature=kwargs.pop('temperature', self.temperature),
             stream=kwargs.pop('stream', streaming),
             max_tokens=kwargs.pop('max_tokens', self.max_tokens),
+            extra_headers=kwargs.pop("extra_headers", self.ext_headers),
+            extra_body=ext_params,
             **kwargs,
         )
         if not streaming:
@@ -105,14 +115,21 @@ class OpenAIStyleLLM(LLM):
         streaming = kwargs.pop("streaming") if "streaming" in kwargs else self.streaming
         if 'stream' in kwargs:
             streaming = kwargs.pop('stream')
-        self.async_client = self._new_async_client()
+        self.async_client = self._new_async_client(kwargs.pop("api_base", None))
         async_client = self.async_client
+        ext_params = self.ext_params.copy()
+        if streaming and "stream_options" not in ext_params:
+            ext_params["stream_options"] = {
+                "include_usage": True
+            }
         chat_completion = await async_client.chat.completions.create(
             messages=messages,
             model=kwargs.pop('model', self.model_name),
             temperature=kwargs.pop('temperature', self.temperature),
             stream=kwargs.pop('stream', streaming),
             max_tokens=kwargs.pop('max_tokens', self.max_tokens),
+            extra_headers=kwargs.pop("extra_headers", self.ext_headers),
+            extra_body=ext_params,
             **kwargs,
         )
         if not streaming:
@@ -147,7 +164,7 @@ class OpenAIStyleLLM(LLM):
         if not isinstance(chunk, dict):
             chunk = chunk.dict()
         if len(chunk["choices"]) == 0:
-            return
+            return LLMOutput(text="", raw=chat_completion.model_dump())
         choice = chunk["choices"][0]
         message = choice.get("delta")
         text = message.get("content")
@@ -186,6 +203,11 @@ class OpenAIStyleLLM(LLM):
         if 'proxy' in component_configer.configer.value:
             proxy = component_configer.configer.value.get('proxy')
             self.proxy = process_yaml_func(proxy, component_configer.yaml_func_instance)
+        if component_configer.configer.value.get("ext_headers"):
+            self.ext_headers = component_configer.configer.value.get("ext_headers")
+        if component_configer.configer.value.get("ext_params"):
+            self.ext_params = component_configer.configer.value.get("ext_params")
+
         return super().initialize_by_component_configer(component_configer)
 
     def get_num_tokens(self, text: str) -> int:

{agentuniverse-0.0.17.dev5 → agentuniverse-0.0.17.dev6}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "agentUniverse"
-version = "0.0.17.dev5"
+version = "0.0.17.dev6"
 description = "agentUniverse is a framework for developing applications powered by multi-agent base on large language model."
 
 authors = ["AntGroup <jerry.zzw@antgroup.com>"]
@@ -64,7 +64,7 @@ jieba = "^0.42.1"
 networkx = "^3.3"
 httpx = "0.27.2"
 tomli = "^2.2"
-mcp = "^1.8.0"
+mcp = "~=1.9.0"
 opentracing = ">=2.4.0,<3.0.0"
 
 [tool.poetry.extras]