Upload 9 files

examples/trace/asyncio_test.py

@@ -0,0 +1,41 @@
+import asyncio
+import aworld.trace as trace
+from aworld.logs.util import logger
+trace.configure()
+
+
+async def async_handler(name):
+    async with trace.span("async_handler") as span:
+        logger.info(f"async_handler start {name}")
+        await asyncio.sleep(1)
+        logger.info(f"async_handler end {name}")
+
+
+async def async_handler2(name):
+    span = trace.get_current_span()
+    logger.info(f"async_handler2 span: {span.get_trace_id()}")
+    logger.info(f"async_handler2 start {name}")
+    await asyncio.sleep(1)
+    logger.info(f"async_handler2 end {name}")
+
+
+async def test1():
+    logger.info(f"hello test1")
+    task = asyncio.create_task(async_handler('test1'))
+    # await task
+    logger.info(f"hello test1 end")
+
+
+async def test2():
+    async with trace.span("test2") as span:
+        logger.info(f"hello test2")
+        task = asyncio.create_task(async_handler2(
+            'test2'))
+        # await task
+        logger.info(f"hello test2 end")
+
+if __name__ == "__main__":
+    with trace.span("hello") as span:
+        logger.info(f"main execute")
+        asyncio.run(test2())
+        asyncio.run(test1())

examples/trace/autotrace_demo.py

@@ -0,0 +1,24 @@
+import time
+
+class TestClassA:
+
+    def classa_function_1(self):
+        print("classa_function_1")
+
+    def classa_function_2(self):
+        time.sleep(0.02)
+        print("classa_function_2")
+
+    def classa_function_3(self):
+        print("classa_function_3")
+
+class TestClassB:
+    def classb_function_1(self):
+        time.sleep(0.02)
+        print("classb_function_1")
+    def classb_function_2(self):
+        a = TestClassA()
+        a.classa_function_1()
+        a.classa_function_2()
+        a.classa_function_3()
+        print("classb_function_2")

examples/trace/instrument_fastapi.py

@@ -0,0 +1,54 @@
+import os
+import time
+import threading
+from aworld.trace.config import ObservabilityConfig
+from aworld.trace.instrumentation.fastapi import instrument_fastapi
+from aworld.trace.instrumentation.requests import instrument_requests
+from aworld.logs.util import logger, trace_logger
+import aworld.trace as trace
+from aworld.utils.import_package import import_packages
+import_packages(['fastapi', 'uvicorn'])  # noqa
+import fastapi  # noqa
+import uvicorn  # noqa
+
+os.environ["MONITOR_SERVICE_NAME"] = "otlp_example"
+os.environ["ANT_OTEL_ENDPOINT"] = "https://antcollector.alipay.com/namespace/aworld/task/aworld/otlp/api/v1/metrics"
+
+trace.configure(ObservabilityConfig(
+    metrics_provider="otlp",
+    metrics_backend="antmonitor"
+))
+
+instrument_fastapi()
+instrument_requests()
+
+app = fastapi.FastAPI()
+
+
+@app.get("/api/hello")
+async def hello():
+    return {"message": "Hello World"}
+
+
+def invoke_api():
+    import requests
+    response = requests.get('http://127.0.0.1:7071/api/hello')
+    logger.info(f"invoke_api response={response.text}")
+
+
+def main():
+    logger.info("main running")
+    with trace.span("test_fastapi") as span:
+        trace_logger.info("start invoke_api")
+        invoke_api()
+
+
+if __name__ == "__main__":
+    server_thread = threading.Thread(
+        target=lambda: uvicorn.run(app, host="0.0.0.0", port=7071),
+        daemon=True
+    )
+    server_thread.start()
+    time.sleep(1)
+    main()
+    server_thread.join()

examples/trace/instrument_flask.py

@@ -0,0 +1,47 @@
+import threading
+import flask
+from aworld.trace.instrumentation.flask import instrument_flask
+from aworld.trace.instrumentation.requests import instrument_requests
+from aworld.logs.util import logger, trace_logger
+import aworld.trace as trace
+import os
+from aworld.trace.config import ObservabilityConfig
+
+os.environ["MONITOR_SERVICE_NAME"] = "otlp_example"
+os.environ["ANT_OTEL_ENDPOINT"] = "https://antcollector.alipay.com/namespace/aworld/task/aworld/otlp/api/v1/metrics"
+
+trace.configure(ObservabilityConfig(
+    metrics_provider="otlp",
+    metrics_backend="antmonitor"
+))
+instrument_flask()
+instrument_requests()
+
+app = flask.Flask(__name__)
+
+
+@app.route('/api/test')
+def test():
+    return 'Hello, World!'
+
+
+thread = threading.Thread(target=lambda: app.run(port=7070), daemon=True)
+thread.start()
+
+
+def invoke_api():
+    import requests
+    response = requests.get('http://localhost:7070/api/test')
+    logger.info(f"invoke_api response={response.text}")
+
+
+def main():
+    logger.info("main running")
+    with trace.span("test_flask") as span:
+        trace_logger.info("start invoke_api")
+        invoke_api()
+
+
+if __name__ == "__main__":
+    main()
+    thread.join()

examples/trace/instrument_threading.py

@@ -0,0 +1,29 @@
+import threading
+import aworld.trace as trace
+import os
+import time
+from aworld.trace.instrumentation.threading import instrument_theading
+from aworld.logs.util import logger, trace_logger
+
+os.environ["MONITOR_SERVICE_NAME"] = "otlp_example"
+trace.configure()
+instrument_theading()
+
+
+def child_thread_func():
+    logger.info("child thread running")
+    with trace.span("child_thread") as span:
+        trace_logger.info("child thread running")
+    time.sleep(1000)
+
+
+def main():
+    logger.info("main running")
+    with trace.span("test_fastapi") as span:
+        trace_logger.info("start run child_thread_func")
+        threading.Thread(target=child_thread_func).start()
+        threading.Thread(target=child_thread_func).start()
+
+
+if __name__ == "__main__":
+    main()

examples/trace/opentelemetry_trace.py

@@ -0,0 +1,69 @@
+import os  # noqa
+# os.environ["START_TRACE_SERVER"] = "false"  # noqa
+os.environ["MONITOR_SERVICE_NAME"] = "otlp_example"  # noqa
+# os.environ["OTLP_TRACES_ENDPOINT"] = "http://localhost:4318/v1/traces"
+# os.environ["METRICS_SYSTEM_ENABLED"] = "true"
+# os.environ["LOGFIRE_WRITE_TOKEN"] = (
+#     "Your logfire write token, "
+#     "create guide refer to "
+#     "https://logfire.pydantic.dev/docs/how-to-guides/create-write-tokens/"
+# )
+
+import aworld.trace as trace  # noqa
+from aworld.logs.util import logger, trace_logger
+from aworld.trace.server import get_trace_server
+
+
+trace.configure()
+
+
+@trace.func_span(span_name="test_func", attributes={"test_attr": "test_value"}, extract_args=["param1"], add_attr="add_attr_value")
+def traced_func(param1: str = None, param2: int = None):
+    trace_logger.info("this is a traced func")
+    traced_func2(param1="func2_param1_value", param2=222)
+    traced_func3(param1="func3_param1_value", param2=333)
+
+
+@trace.func_span(span_name="test_func_2", add_attr="add_attr_value")
+def traced_func2(param1: str = None, param2: int = None):
+    name = 'func2'
+    trace_logger.info(f"this is a traced {name}")
+    raise Exception("this is a traced func2 exception")
+
+
+@trace.func_span
+def traced_func3(param1: str = None, param2: int = None):
+    trace_logger.info("this is a traced func3")
+
+
+def main():
+
+    logger.info("this is a no trace log")
+
+    trace.auto_tracing("examples.trace.*", 0.01)
+
+    with trace.span("hello") as span:
+        span.set_attribute("parent_test_attr", "pppppp")
+        logger.info("hello aworld")
+        trace_logger.info("trace hello aworld")
+        with trace.span("child hello") as span2:
+            span2.set_attribute("child_test_attr", "cccccc")
+            logger.info("child hello aworld")
+            current_span = trace.get_current_span()
+            logger.info("trace_id=%s", current_span.get_trace_id())
+    try:
+        traced_func(param1="func1_param1_value", param2=111)
+    except Exception as e:
+        logger.error(f"exception: {e}")
+    # from examples.trace.autotrace_demo import TestClassB
+    # b = TestClassB()
+    # b.classb_function_1()
+    # b.classb_function_2()
+    # b.classb_function_1()
+    # b.classb_function_2()
+    if get_trace_server():
+        get_trace_server().join()
+
+
+if __name__ == "__main__":
+    main()

examples/trace/sofa_tracer.py

@@ -0,0 +1,71 @@
+import os  # noqa: E402
+
+os.environ["MONITOR_SERVICE_NAME"] = "otlp_example"  # noqa
+os.environ["ANT_OTEL_ENDPOINT"] = "https://antcollector.alipay.com/namespace/aworld/task/aworld/otlp/api/v1/metrics"  # noqa
+os.environ["OTLP_TRACES_ENDPOINT"] = "https://antcollector.alipay.com/namespace/aworld/task/aworld_trace/otlp/api/v1/traces"  # noqa
+
+from aworld.trace.config import ObservabilityConfig
+from aworld.logs.util import logger
+from aworld.trace.baggage import BaggageContext
+from aworld.trace.base import get_tracer_provider
+from aworld.trace.instrumentation.requests import instrument_requests
+from aworld.trace.instrumentation.flask import instrument_flask
+import flask
+import threading
+import aworld.trace as trace
+
+
+trace.configure(ObservabilityConfig(
+    trace_provider="otlp",
+    trace_backends=["other_otlp"],
+    trace_base_url="https://antcollector.alipay.com/namespace/aworld/task/aworld_trace/otlp/api/v1/traces",
+    metrics_provider="otlp",
+    metrics_backend="antmonitor",
+    metrics_base_url="https://antcollector.alipay.com/namespace/aworld/task/aworld/otlp/api/v1/metrics"
+))
+instrument_flask()
+instrument_requests()
+
+app = flask.Flask(__name__)
+
+
+@app.route('/api/test')
+def test():
+    sofa_trace_id = BaggageContext.get_baggage_value("attributes.sofa.traceid")
+    sofa_rpc_id = BaggageContext.get_baggage_value("attributes.sofa.rpcid")
+    sofa_pen_attrs = BaggageContext.get_baggage_value(
+        "attributes.sofa.penattrs")
+    sofa_sys_pen_attrs = BaggageContext.get_baggage_value(
+        "attributes.sofa.syspenattrs")
+    logger.info(
+        f"test sofa_trace_id={sofa_trace_id}, sofa_rpc_id={sofa_rpc_id}, sofa_pen_attrs={sofa_pen_attrs}, sofa_sys_pen_attrs={sofa_sys_pen_attrs}"
+    )
+    return 'Hello, World!'
+
+
+thread = threading.Thread(target=lambda: app.run(port=7070), daemon=True)
+thread.start()
+
+
+def invoke_api():
+    import requests
+    session = requests.session()
+    session.headers.update({
+        "SOFA-TraceId": "12345678901234567890123456789012",
+        "SOFA-RpcId": "0.1.1",
+        "sofaPenAttrs": "key1=value1&key2=value2",
+        "sysPenAttrs": "key1=value1&key2=value2"
+    })
+    response = session.get('http://localhost:7070/api/test')
+    logger.info(f"invoke_api response={response.text}")
+
+
+def main():
+    logger.info("main running")
+    invoke_api()
+
+
+if __name__ == "__main__":
+    main()
+    get_tracer_provider().force_flush(1000)
+    thread.join()

examples/trace/span_cosumer.py

@@ -0,0 +1,33 @@
+import os
+import json
+from aworld.logs.util import logger, trace_logger
+from typing import Sequence
+import aworld.trace as trace
+from aworld.trace.base import Span
+from aworld.trace.span_cosumer import register_span_consumer, SpanConsumer
+from aworld.logs.util import logger, trace_logger
+
+os.environ["MONITOR_SERVICE_NAME"] = "otlp_example"
+
+
+@register_span_consumer({"test_param": "MockSpanConsumer111"})
+class MockSpanConsumer(SpanConsumer):
+
+    def __init__(self, test_param=None):
+        self._test_param = test_param
+
+    def consume(self, spans: Sequence[Span]) -> None:
+        for span in spans:
+            logger.info(
+                f"_test_param={self._test_param}, trace_id={span.get_trace_id()}, span_id={span.get_span_id()}, attributes={span.attributes}")
+
+
+def main():
+    with trace.span("hello") as span:
+        span.set_attribute("parent_test_attr", "pppppp")
+        logger.info("hello aworld")
+        trace_logger.info("trace hello aworld")
+
+
+if __name__ == "__main__":
+    main()

examples/trace/trace_agent.py

@@ -0,0 +1,249 @@
+import traceback
+from aworld.agents.llm_agent import Agent
+from aworld.config.conf import AgentConfig, ConfigDict
+from aworld.core.common import Observation, ActionModel
+from typing import Dict, Any, List, Union, Callable
+from aworld.core.tool.base import ToolFactory
+from aworld.models.llm import call_llm_model, acall_llm_model
+from aworld.utils.common import sync_exec
+from aworld.logs.util import logger
+from examples.tools.common import Tools
+from examples.tools.tool_action import GetTraceAction
+from aworld.core.agent.swarm import Swarm
+from aworld.runner import Runners
+from aworld.trace.server import get_trace_server
+from aworld.runners.state_manager import RuntimeStateManager, RunNode
+import aworld.trace as trace
+
+
+trace.configure()
+
+
+class TraceAgent(Agent):
+
+    def __init__(self,
+                 conf: Union[Dict[str, Any], ConfigDict, AgentConfig],
+                 resp_parse_func: Callable[..., Any] = None,
+                 **kwargs):
+        super().__init__(conf, **kwargs)
+
+    def policy(self, observation: Observation, info: Dict[str, Any] = {}, **kwargs) -> List[ActionModel]:
+        """use trace tool to get trace data, and call llm to summary
+
+        Args:
+            observation: The state observed from tools in the environment.
+            info: Extended information is used to assist the agent to decide a policy.
+
+        Returns:
+            ActionModel sequence from agent policy
+        """
+
+        self._finished = False
+        self.desc_transform()
+
+        tool_name = "trace"
+        tool = ToolFactory(tool_name, asyn=False)
+        tool.reset()
+        tool_params = {}
+        action = ActionModel(tool_name=tool_name,
+                             action_name=GetTraceAction.GET_TRACE.name,
+                             agent_name=self.id(),
+                             params=tool_params)
+        message = tool.step(action)
+
+        observation, _, _, _, _ = message.payload
+
+        llm_response = None
+
+        messages = self.messages_transform(content=observation.content,
+                                           sys_prompt=self.system_prompt,
+                                           agent_prompt=self.agent_prompt)
+        try:
+            llm_response = call_llm_model(
+                self.llm,
+                messages=messages,
+                model=self.model_name,
+                temperature=self.conf.llm_config.llm_temperature
+            )
+
+            logger.info(f"Execute response: {llm_response.message}")
+        except Exception as e:
+            logger.warn(traceback.format_exc())
+            raise e
+        finally:
+            if llm_response:
+                if llm_response.error:
+                    logger.info(
+                        f"{self.id()} llm result error: {llm_response.error}")
+            else:
+                logger.error(f"{self.id()} failed to get LLM response")
+                raise RuntimeError(
+                    f"{self.id()} failed to get LLM response")
+
+        agent_result = sync_exec(self.resp_parse_func, llm_response)
+        if not agent_result.is_call_tool:
+            self._finished = True
+        return agent_result.actions
+
+    async def async_policy(self, observation: Observation, info: Dict[str, Any] = {}, **kwargs) -> List[ActionModel]:
+
+        self._finished = False
+        self.desc_transform()
+
+        tool_name = "trace"
+        tool = ToolFactory(tool_name, asyn=False)
+        tool.reset()
+        tool_params = {}
+        action = ActionModel(tool_name=tool_name,
+                             action_name=GetTraceAction.GET_TRACE.name,
+                             agent_name=self.id(),
+                             params=tool_params)
+        message = tool.step([action])
+
+        observation, _, _, _, _ = message.payload
+
+        llm_response = None
+
+        messages = self.messages_transform(content=observation.content,
+                                           sys_prompt=self.system_prompt,
+                                           agent_prompt=self.agent_prompt)
+        try:
+            llm_response = await acall_llm_model(
+                self.llm,
+                messages=messages,
+                model=self.model_name,
+                temperature=self.conf.llm_config.llm_temperature
+            )
+
+            logger.info(f"Execute response: {llm_response.message}")
+        except Exception as e:
+            logger.warn(traceback.format_exc())
+            raise e
+        finally:
+            if llm_response:
+                if llm_response.error:
+                    logger.info(
+                        f"{self.id()} llm result error: {llm_response.error}")
+            else:
+                logger.error(f"{self.id()} failed to get LLM response")
+                raise RuntimeError(
+                    f"{self.id()} failed to get LLM response")
+
+        agent_result = sync_exec(self.resp_parse_func, llm_response)
+        if not agent_result.is_call_tool:
+            self._finished = True
+        return agent_result.actions
+
+
+search_sys_prompt = "You are a helpful search agent."
+search_prompt = """
+    Please act as a search agent, constructing appropriate keywords and searach terms, using search toolkit to collect relevant information, including urls, webpage snapshots, etc.
+
+    Here are the question: {task}
+
+    pleas only use one action complete this task, at least results 6 pages.
+    """
+
+summary_sys_prompt = "You are a helpful general summary agent."
+
+summary_prompt = """
+Summarize the following text in one clear and concise paragraph, capturing the key ideas without missing critical points. 
+Ensure the summary is easy to understand and avoids excessive detail.
+
+Here are the content: 
+{task}
+"""
+
+trace_sys_prompt = "You are a helpful trace agent."
+
+trace_prompt = """
+    Please act as a trace agent, Using the provided trace data, summarize the token usage of each agent,
+    whether the runotype attribute of span is an agent or a large model call: 
+        run_type=AGNET represents the agent, 
+        run_type=LLM represents the large model call.
+    The LLM call of a certain agent is represented as LLM span, which is a child span of that agent span
+
+    Here are the content: {task}
+    """
+
+
+def build_run_flow(nodes: List[RunNode]):
+    graph = {}
+    start_nodes = []
+
+    for node in nodes:
+        if hasattr(node, 'parent_node_id') and node.parent_node_id:
+            if node.parent_node_id not in graph:
+                graph[node.parent_node_id] = []
+            graph[node.parent_node_id].append(node.node_id)
+        else:
+            start_nodes.append(node.node_id)
+
+    for start in start_nodes:
+        print("-----------------------------------")
+        _print_tree(graph, start, "", True)
+        print("-----------------------------------")
+
+
+def _print_tree(graph, node_id, prefix, is_last):
+    print(prefix + ("└── " if is_last else "├── ") + node_id)
+    if node_id in graph:
+        children = graph[node_id]
+        for i, child in enumerate(children):
+            _print_tree(graph, child, prefix +
+                        ("    " if is_last else "│   "), i == len(children)-1)
+
+
+if __name__ == "__main__":
+    agent_config = AgentConfig(
+        llm_provider="openai",
+        llm_model_name="gpt-4o",
+        llm_temperature=0.3,
+
+        llm_base_url="http://localhost:34567",
+        llm_api_key="dummy-key",
+    )
+
+    search = Agent(
+        conf=agent_config,
+        name="search_agent",
+        system_prompt=search_sys_prompt,
+        agent_prompt=search_prompt,
+        tool_names=[Tools.SEARCH_API.value]
+    )
+
+    summary = Agent(
+        conf=agent_config,
+        name="summary_agent",
+        system_prompt=summary_sys_prompt,
+        agent_prompt=summary_prompt
+    )
+
+    trace = TraceAgent(
+        conf=agent_config,
+        name="trace_agent",
+        system_prompt=trace_sys_prompt,
+        agent_prompt=trace_prompt
+    )
+
+    # default is sequence swarm mode
+    swarm = Swarm(search, summary, trace, max_steps=1, event_driven=True)
+
+    prefix = "search baidu:"
+    # can special search google, wiki, duck go, or baidu. such as:
+    # prefix = "search wiki: "
+    try:
+        res = Runners.sync_run(
+            input=prefix + """What is an agent.""",
+            swarm=swarm,
+            session_id="123"
+        )
+        print(res.answer)
+    except Exception as e:
+        logger.error(traceback.format_exc())
+
+    state_manager = RuntimeStateManager.instance()
+    nodes = state_manager.get_nodes("123")
+    logger.info(f"session 123 nodes: {nodes}")
+    build_run_flow(nodes)
+    get_trace_server().join()