diff --git a/agentops/__init__.py b/agentops/__init__.py index 29557e7d0..947ddd530 100755 --- a/agentops/__init__.py +++ b/agentops/__init__.py @@ -61,6 +61,7 @@ def init( log_level: Optional[Union[str, int]] = None, fail_safe: Optional[bool] = None, exporter_endpoint: Optional[str] = None, + session_name: Optional[str] = None, **kwargs, ): """ @@ -88,6 +89,7 @@ def init( fail_safe (bool): Whether to suppress errors and continue execution when possible. exporter_endpoint (str, optional): Endpoint for the exporter. If none is provided, key will be read from the AGENTOPS_EXPORTER_ENDPOINT environment variable. + session_name (str, optional): Name of the session to be used in the span attributes. **kwargs: Additional configuration parameters to be passed to the client. """ global _client @@ -116,6 +118,7 @@ def init( log_level=log_level, fail_safe=fail_safe, exporter_endpoint=exporter_endpoint, + session_name=session_name, **kwargs, ) diff --git a/agentops/client/client.py b/agentops/client/client.py index 9f29dcc92..7d4b25b03 100644 --- a/agentops/client/client.py +++ b/agentops/client/client.py @@ -99,10 +99,18 @@ def init(self, **kwargs): if self.config.auto_start_session: from agentops.legacy import start_session - # Pass default_tags if they exist - if self.config.default_tags: + if self.config.default_tags and self.config.session_name: + logger.debug(f"Starting session with tags: {self.config.default_tags}") + logger.debug(f"Starting session with name: {self.config.session_name}") + session = start_session(session_name=self.config.session_name, tags=list(self.config.default_tags)) + elif self.config.default_tags: + logger.debug(f"Starting session with tags: {self.config.default_tags}") session = start_session(tags=list(self.config.default_tags)) + elif self.config.session_name: + logger.debug(f"Starting session with name: {self.config.session_name}") + session = start_session(session_name=self.config.session_name) else: + logger.debug("Starting session without tags or name") session = start_session() # Register this session globally diff --git a/agentops/config.py b/agentops/config.py index 6af2005c4..d931f404e 100644 --- a/agentops/config.py +++ b/agentops/config.py @@ -124,6 +124,11 @@ class Config: default_factory=lambda: None, metadata={"description": "Custom span processor for OpenTelemetry trace data"} ) + session_name: Optional[str] = field( + default_factory=lambda: None, + metadata={"description": "Name of the session to be used in the span attributes"}, + ) + def configure( self, api_key: Optional[str] = None, @@ -144,6 +149,7 @@ def configure( exporter: Optional[SpanExporter] = None, processor: Optional[SpanProcessor] = None, exporter_endpoint: Optional[str] = None, + session_name: Optional[str] = None, ): """Configure settings from kwargs, validating where necessary""" if api_key is not None: @@ -214,6 +220,9 @@ def configure( # else: # self.exporter_endpoint = self.endpoint + if session_name is not None: + self.session_name = session_name + def dict(self): """Return a dictionary representation of the config""" return { @@ -235,6 +244,7 @@ def dict(self): "exporter": self.exporter, "processor": self.processor, "exporter_endpoint": self.exporter_endpoint, + "session_name": self.session_name, } def json(self): diff --git a/agentops/instrumentation/crewai/instrumentation.py b/agentops/instrumentation/crewai/instrumentation.py index 3a2964733..afb5eff02 100644 --- a/agentops/instrumentation/crewai/instrumentation.py +++ b/agentops/instrumentation/crewai/instrumentation.py @@ -14,6 +14,8 @@ from agentops.instrumentation.crewai.version import __version__ from agentops.semconv import SpanAttributes, AgentOpsSpanKindValues, Meters, ToolAttributes, MessageAttributes from .crewai_span_attributes import CrewAISpanAttributes, set_span_attribute +from agentops import get_client + # Initialize logger logger = logging.getLogger(__name__) @@ -156,15 +158,26 @@ def wrap_kickoff( args, kwargs, ): + span_name = "crewai.workflow" logger.debug( f"CrewAI: Starting workflow instrumentation for Crew with {len(getattr(instance, 'agents', []))} agents" ) + config = get_client().config + attributes = { + SpanAttributes.LLM_SYSTEM: "crewai", + } + + if config.session_name: + span_name = f"{config.session_name}.workflow" + + if config.default_tags and len(config.default_tags) > 0: + tag_list = list(config.default_tags) + attributes[SpanAttributes.AGENTOPS_SPAN_TAGS] = tag_list + with tracer.start_as_current_span( - "crewai.workflow", + name=span_name, kind=SpanKind.INTERNAL, - attributes={ - SpanAttributes.LLM_SYSTEM: "crewai", - }, + attributes=attributes, ) as span: try: span.set_attribute(TELEMETRY_SDK_NAME, "agentops") @@ -327,12 +340,19 @@ def wrap_agent_execute_task( tracer, duration_histogram, token_histogram, environment, application_name, wrapped, instance, args, kwargs ): agent_name = instance.role if hasattr(instance, "role") else "agent" + attributes = { + SpanAttributes.AGENTOPS_SPAN_KIND: AgentOpsSpanKindValues.AGENT.value, + } + + config = get_client().config + if config.default_tags and len(config.default_tags) > 0: + tag_list = list(config.default_tags) + attributes[SpanAttributes.AGENTOPS_SPAN_TAGS] = tag_list + with tracer.start_as_current_span( f"{agent_name}.agent", kind=SpanKind.CLIENT, - attributes={ - SpanAttributes.AGENTOPS_SPAN_KIND: AgentOpsSpanKindValues.AGENT.value, - }, + attributes=attributes, ) as span: try: span.set_attribute(TELEMETRY_SDK_NAME, "agentops") @@ -381,12 +401,22 @@ def wrap_task_execute( ): task_name = instance.description if hasattr(instance, "description") else "task" + config = get_client().config + attributes = { + SpanAttributes.AGENTOPS_SPAN_KIND: AgentOpsSpanKindValues.TASK.value, + } + + if config.default_tags and len(config.default_tags) > 0: + tag_list = list(config.default_tags) + attributes[SpanAttributes.AGENTOPS_SPAN_TAGS] = tag_list + + if config.session_name: + task_name = config.session_name + with tracer.start_as_current_span( f"{task_name}.task", kind=SpanKind.CLIENT, - attributes={ - SpanAttributes.AGENTOPS_SPAN_KIND: AgentOpsSpanKindValues.TASK.value, - }, + attributes=attributes, ) as span: try: span.set_attribute(TELEMETRY_SDK_NAME, "agentops") @@ -411,7 +441,16 @@ def wrap_llm_call( tracer, duration_histogram, token_histogram, environment, application_name, wrapped, instance, args, kwargs ): llm = instance.model if hasattr(instance, "model") else "llm" - with tracer.start_as_current_span(f"{llm}.llm", kind=SpanKind.CLIENT, attributes={}) as span: + attributes = { + SpanAttributes.AGENTOPS_SPAN_KIND: AgentOpsSpanKindValues.LLM.value, + } + + config = get_client().config + if config.default_tags and len(config.default_tags) > 0: + tag_list = list(config.default_tags) + attributes[SpanAttributes.AGENTOPS_SPAN_TAGS] = tag_list + + with tracer.start_as_current_span(f"{llm}.llm", kind=SpanKind.CLIENT, attributes=attributes) as span: start_time = time.time() try: span.set_attribute(TELEMETRY_SDK_NAME, "agentops") diff --git a/agentops/legacy/__init__.py b/agentops/legacy/__init__.py index f733f2aeb..780e0318a 100644 --- a/agentops/legacy/__init__.py +++ b/agentops/legacy/__init__.py @@ -73,7 +73,9 @@ def end_session(self, **kwargs): _flush_span_processors() -def _create_session_span(tags: Union[Dict[str, Any], List[str], None] = None) -> tuple: +def _create_session_span( + session_name: Optional[str] = None, tags: Union[Dict[str, Any], List[str], None] = None +) -> tuple: """ Helper function to create a session span with tags. @@ -94,10 +96,13 @@ def _create_session_span(tags: Union[Dict[str, Any], List[str], None] = None) -> attributes = {} if tags: attributes["tags"] = tags - return _make_span("session", span_kind=SpanKind.SESSION, attributes=attributes) + + span_name = "session" if session_name is None else session_name + return _make_span(span_name, span_kind=SpanKind.SESSION, attributes=attributes) def start_session( + session_name: Optional[str] = None, tags: Union[Dict[str, Any], List[str], None] = None, ) -> Session: """ @@ -153,7 +158,7 @@ def start_session( _current_session = dummy_session return dummy_session - span, ctx, token = _create_session_span(tags) + span, ctx, token = _create_session_span(session_name, tags) session = Session(span, token) # Set the global session reference diff --git a/agentops/sdk/decorators/__init__.py b/agentops/sdk/decorators/__init__.py index 706bd4624..c17b08aa7 100644 --- a/agentops/sdk/decorators/__init__.py +++ b/agentops/sdk/decorators/__init__.py @@ -14,8 +14,9 @@ operation = create_entity_decorator(SpanKind.OPERATION) workflow = create_entity_decorator(SpanKind.WORKFLOW) session = create_entity_decorator(SpanKind.SESSION) +tool = create_entity_decorator(SpanKind.TOOL) operation = task -__all__ = ["agent", "task", "workflow", "session", "operation"] +__all__ = ["agent", "task", "workflow", "session", "operation", "tool"] # Create decorators task, workflow, session, agent diff --git a/agentops/sdk/decorators/factory.py b/agentops/sdk/decorators/factory.py index 13b09789d..582005ccd 100644 --- a/agentops/sdk/decorators/factory.py +++ b/agentops/sdk/decorators/factory.py @@ -2,10 +2,12 @@ import functools import asyncio + import wrapt # type: ignore from agentops.logging import logger from agentops.sdk.core import TracingCore +from agentops.semconv.span_attributes import SpanAttributes from .utility import ( _create_as_current_span, @@ -28,10 +30,10 @@ def create_entity_decorator(entity_kind: str): A decorator with optional arguments for name and version """ - def decorator(wrapped=None, *, name=None, version=None): + def decorator(wrapped=None, *, name=None, version=None, cost=None): # Handle case where decorator is called with parameters if wrapped is None: - return functools.partial(decorator, name=name, version=version) + return functools.partial(decorator, name=name, version=version, cost=cost) # Handle class decoration if inspect.isclass(wrapped): @@ -39,6 +41,7 @@ def decorator(wrapped=None, *, name=None, version=None): class WrappedClass(wrapped): def __init__(self, *args, **kwargs): operation_name = name or wrapped.__name__ + self._agentops_span_context_manager = _create_as_current_span(operation_name, entity_kind, version) self._agentops_active_span = self._agentops_span_context_manager.__enter__() @@ -51,23 +54,18 @@ def __init__(self, *args, **kwargs): super().__init__(*args, **kwargs) async def __aenter__(self): - # Added for async context manager support - # This allows using the class with 'async with' statement - # If span is already created in __init__, just return self if hasattr(self, "_agentops_active_span") and self._agentops_active_span is not None: return self # Otherwise create span (for backward compatibility) operation_name = name or wrapped.__name__ + self._agentops_span_context_manager = _create_as_current_span(operation_name, entity_kind, version) self._agentops_active_span = self._agentops_span_context_manager.__enter__() return self async def __aexit__(self, exc_type, exc_val, exc_tb): - # Added for proper async cleanup - # This ensures spans are properly closed when using 'async with' - if hasattr(self, "_agentops_active_span") and hasattr(self, "_agentops_span_context_manager"): try: _record_entity_output(self._agentops_active_span, self) @@ -104,10 +102,12 @@ def wrapper(wrapped, instance, args, kwargs): # Handle generator functions if is_generator: - # Use the old approach for generators span, ctx, token = _make_span(operation_name, entity_kind, version) try: _record_entity_input(span, args, kwargs) + # Set cost attribute if tool + if entity_kind == "tool" and cost is not None: + span.set_attribute(SpanAttributes.LLM_USAGE_TOOL_COST, cost) except Exception as e: logger.warning(f"Failed to record entity input: {e}") @@ -116,10 +116,12 @@ def wrapper(wrapped, instance, args, kwargs): # Handle async generator functions elif is_async_generator: - # Use the old approach for async generators span, ctx, token = _make_span(operation_name, entity_kind, version) try: _record_entity_input(span, args, kwargs) + # Set cost attribute if tool + if entity_kind == "tool" and cost is not None: + span.set_attribute(SpanAttributes.LLM_USAGE_TOOL_COST, cost) except Exception as e: logger.warning(f"Failed to record entity input: {e}") @@ -133,6 +135,9 @@ async def _wrapped_async(): with _create_as_current_span(operation_name, entity_kind, version) as span: try: _record_entity_input(span, args, kwargs) + # Set cost attribute if tool + if entity_kind == "tool" and cost is not None: + span.set_attribute(SpanAttributes.LLM_USAGE_TOOL_COST, cost) except Exception as e: logger.warning(f"Failed to record entity input: {e}") @@ -144,6 +149,7 @@ async def _wrapped_async(): logger.warning(f"Failed to record entity output: {e}") return result except Exception as e: + logger.error(f"Error in async function execution: {e}") span.record_exception(e) raise @@ -154,7 +160,9 @@ async def _wrapped_async(): with _create_as_current_span(operation_name, entity_kind, version) as span: try: _record_entity_input(span, args, kwargs) - + # Set cost attribute if tool + if entity_kind == "tool" and cost is not None: + span.set_attribute(SpanAttributes.LLM_USAGE_TOOL_COST, cost) except Exception as e: logger.warning(f"Failed to record entity input: {e}") @@ -167,6 +175,7 @@ async def _wrapped_async(): logger.warning(f"Failed to record entity output: {e}") return result except Exception as e: + logger.error(f"Error in sync function execution: {e}") span.record_exception(e) raise diff --git a/agentops/semconv/span_attributes.py b/agentops/semconv/span_attributes.py index 79f0285a9..0931b8d64 100644 --- a/agentops/semconv/span_attributes.py +++ b/agentops/semconv/span_attributes.py @@ -65,6 +65,7 @@ class SpanAttributes: LLM_USAGE_CACHE_READ_INPUT_TOKENS = "gen_ai.usage.cache_read_input_tokens" LLM_USAGE_REASONING_TOKENS = "gen_ai.usage.reasoning_tokens" LLM_USAGE_STREAMING_TOKENS = "gen_ai.usage.streaming_tokens" + LLM_USAGE_TOOL_COST = "gen_ai.usage.total_cost" # Message attributes # see ./message.py for message-related attributes @@ -86,6 +87,7 @@ class SpanAttributes: AGENTOPS_ENTITY_INPUT = "agentops.entity.input" AGENTOPS_SPAN_KIND = "agentops.span.kind" AGENTOPS_ENTITY_NAME = "agentops.entity.name" + AGENTOPS_SPAN_TAGS = "tags" # Operation attributes OPERATION_NAME = "operation.name" diff --git a/docs/v1/concepts/sessions.mdx b/docs/v1/concepts/sessions.mdx index ec98cb7f5..d9996834a 100644 --- a/docs/v1/concepts/sessions.mdx +++ b/docs/v1/concepts/sessions.mdx @@ -27,6 +27,7 @@ Optionally, sessions may include: - **Tags**: Tags allow for the categorization and later retrieval of sessions. - **Host Environment**: Automatically gathers basic information about the system on which the session ran. - **Video**: If applicable, an optional video recording of the session. +- **Session Name**: A custom name for the session that helps identify and organize different types of sessions in the dashboard. Can be set during initialization or when starting a session. ### Methods #### `end_session` diff --git a/docs/v1/usage/sdk-reference.mdx b/docs/v1/usage/sdk-reference.mdx index ae0690c3d..f1f565646 100644 --- a/docs/v1/usage/sdk-reference.mdx +++ b/docs/v1/usage/sdk-reference.mdx @@ -26,6 +26,7 @@ The first element of AgentOps is always calling .init() - `auto_start_session` (bool): Whether to start a session automatically when the client is created. You may wish to delay starting a session in order to do additional setup or starting a session on a child process. - `inherited_session_id` (str, optional): When creating the client, passing in this value will connect the client to an existing session. This is useful when having separate processes contribute to the same session. - `skip_auto_end_session` (bool, optional): If you are using a framework such as Crew, the framework can decide when to halt execution. Setting this parameter to true will not end your agentops session when this happens. +- `session_name` (str, optional): Name of the session to be used in the span attributes. This helps identify and organize different types of sessions in the dashboard. **Returns**: diff --git a/docs/v2/concepts/decorators.mdx b/docs/v2/concepts/decorators.mdx index 36206a94e..4596e4193 100644 --- a/docs/v2/concepts/decorators.mdx +++ b/docs/v2/concepts/decorators.mdx @@ -13,6 +13,7 @@ AgentOps provides the following decorators: | `@operation` | Track discrete operations performed by agents | OPERATION span | | `@workflow` | Track a sequence of operations | WORKFLOW span | | `@task` | Track smaller units of work (similar to operations) | TASK span | +| `@tool` | Track tool usage and cost in agent operations | TOOL span | ## Decorator Hierarchy @@ -190,6 +191,50 @@ class DataProcessor: The `@task` and `@operation` decorators function identically (they are aliases in the codebase), and you can choose the one that best fits your semantic needs. +### @tool + +The `@tool` decorator tracks tool usage within agent operations and supports cost tracking. It works with all function types: synchronous, asynchronous, generator, and async generator. + +```python +from agentops.sdk.decorators import agent, tool +import asyncio + +@agent +class ProcessingAgent: + def __init__(self): + pass + + @tool(cost=0.01) + def sync_tool(self, item): + """Synchronous tool with cost tracking.""" + return f"Processed {item}" + + @tool(cost=0.02) + async def async_tool(self, item): + """Asynchronous tool with cost tracking.""" + await asyncio.sleep(0.1) + return f"Async processed {item}" + + @tool(cost=0.03) + def generator_tool(self, items): + """Generator tool with cost tracking.""" + for item in items: + yield self.sync_tool(item) + + @tool(cost=0.04) + async def async_generator_tool(self, items): + """Async generator tool with cost tracking.""" + for item in items: + await asyncio.sleep(0.1) + yield await self.async_tool(item) +``` + +The tool decorator provides: +- Cost tracking for each tool call +- Proper span creation and nesting +- Support for all function types (sync, async, generator, async generator) +- Cost accumulation in generator and async generator operations + ## Decorator Attributes You can pass additional attributes to decorators: diff --git a/docs/v2/concepts/sessions.mdx b/docs/v2/concepts/sessions.mdx index b2eae78f2..5f9b78df1 100644 --- a/docs/v2/concepts/sessions.mdx +++ b/docs/v2/concepts/sessions.mdx @@ -13,6 +13,19 @@ import agentops agentops.init(api_key="YOUR_API_KEY", tags=["production"]) ``` +You can also provide a custom name for your session during initialization: + +```python +import agentops + +# Initialize with a custom session name +agentops.init( + api_key="YOUR_API_KEY", + session_name="customer-support-bot", + tags=["production"] +) +``` + This approach: - Creates a session automatically when you initialize the SDK - Tracks all events in the context of this session diff --git a/docs/v2/concepts/tags.mdx b/docs/v2/concepts/tags.mdx index 40ea2bbd1..42b478698 100644 --- a/docs/v2/concepts/tags.mdx +++ b/docs/v2/concepts/tags.mdx @@ -13,6 +13,7 @@ import agentops agentops.init( api_key="YOUR_API_KEY", tags=["production", "customer-service", "gpt-4"] + session_name="customer-support-agent" ) ``` diff --git a/docs/v2/usage/sdk-reference.mdx b/docs/v2/usage/sdk-reference.mdx index 7bd33a45b..4f796c392 100644 --- a/docs/v2/usage/sdk-reference.mdx +++ b/docs/v2/usage/sdk-reference.mdx @@ -33,6 +33,7 @@ Initializes the AgentOps SDK and automatically starts tracking your application. - `fail_safe` (bool, optional): Whether to suppress errors and continue execution when possible. Defaults to False. - `exporter_endpoint` (str, optional): Endpoint for the exporter. If not provided, will be read from the `AGENTOPS_EXPORTER_ENDPOINT` environment variable. Defaults to 'https://otlp.agentops.ai/v1/traces'. - `export_flush_interval` (int, optional): Time interval in milliseconds between automatic exports of telemetry data. Defaults to 1000. +- `session_name` (str, optional): Name of the session to be used in the span attributes. This helps identify and organize different types of sessions in the dashboard. **Returns**: @@ -145,6 +146,7 @@ my_workflow() - `@agent`: Creates an agent span for tracking agent operations - `@operation` / `@task`: Creates operation/task spans for tracking specific operations (these are aliases) - `@workflow`: Creates workflow spans for organizing related operations +- `@tool`: Creates tool spans for tracking tool usage and cost in agent operations. Supports cost parameter for tracking tool usage costs. See [Decorators](/v2/concepts/decorators) for more detailed documentation on using these decorators. diff --git a/tests/unit/sdk/test_decorators.py b/tests/unit/sdk/test_decorators.py index 98f1cd402..80e2e7ddf 100644 --- a/tests/unit/sdk/test_decorators.py +++ b/tests/unit/sdk/test_decorators.py @@ -2,7 +2,7 @@ import asyncio import pytest -from agentops.sdk.decorators import agent, operation, session, workflow, task +from agentops.sdk.decorators import agent, operation, session, workflow, task, tool from agentops.semconv import SpanKind from agentops.semconv.span_attributes import SpanAttributes from tests.unit.sdk.instrumentation_tester import InstrumentationTester @@ -624,3 +624,140 @@ def __init__(self): with pytest.raises(ValueError): async with TestClass() as instance: raise ValueError("Trigger exception for __aexit__ coverage") + + +class TestToolDecorator: + """Tests for the tool decorator functionality.""" + + @pytest.fixture + def agent_class(self): + @agent + class TestAgent: + @tool(cost=0.01) + def process_item(self, item): + return f"Processed {item}" + + @tool(cost=0.02) + async def async_process_item(self, item): + await asyncio.sleep(0.1) + return f"Async processed {item}" + + @tool(cost=0.03) + def generator_process_items(self, items): + for item in items: + yield self.process_item(item) + + @tool(cost=0.04) + async def async_generator_process_items(self, items): + for item in items: + await asyncio.sleep(0.1) + yield await self.async_process_item(item) + + return TestAgent() + + def test_sync_tool_cost(self, agent_class, instrumentation: InstrumentationTester): + """Test synchronous tool with cost attribute.""" + result = agent_class.process_item("test") + + assert result == "Processed test" + + spans = instrumentation.get_finished_spans() + tool_span = next( + span for span in spans if span.attributes.get(SpanAttributes.AGENTOPS_SPAN_KIND) == SpanKind.TOOL + ) + assert tool_span.attributes.get(SpanAttributes.LLM_USAGE_TOOL_COST) == 0.01 + + @pytest.mark.asyncio + async def test_async_tool_cost(self, agent_class, instrumentation: InstrumentationTester): + """Test asynchronous tool with cost attribute.""" + result = await agent_class.async_process_item("test") + + assert result == "Async processed test" + + spans = instrumentation.get_finished_spans() + tool_span = next( + span for span in spans if span.attributes.get(SpanAttributes.AGENTOPS_SPAN_KIND) == SpanKind.TOOL + ) + assert tool_span.attributes.get(SpanAttributes.LLM_USAGE_TOOL_COST) == 0.02 + + def test_generator_tool_cost(self, agent_class, instrumentation: InstrumentationTester): + """Test generator tool with cost attribute.""" + items = ["item1", "item2", "item3"] + results = list(agent_class.generator_process_items(items)) + + assert len(results) == 3 + assert results[0] == "Processed item1" + assert results[1] == "Processed item2" + assert results[2] == "Processed item3" + + spans = instrumentation.get_finished_spans() + tool_spans = [span for span in spans if span.attributes.get(SpanAttributes.AGENTOPS_SPAN_KIND) == SpanKind.TOOL] + assert len(tool_spans) == 4 # Only one span for the generator + assert tool_spans[0].attributes.get(SpanAttributes.LLM_USAGE_TOOL_COST) == 0.01 + assert tool_spans[3].attributes.get(SpanAttributes.LLM_USAGE_TOOL_COST) == 0.03 + + @pytest.mark.asyncio + async def test_async_generator_tool_cost(self, agent_class, instrumentation: InstrumentationTester): + """Test async generator tool with cost attribute.""" + items = ["item1", "item2", "item3"] + results = [result async for result in agent_class.async_generator_process_items(items)] + + assert len(results) == 3 + assert results[0] == "Async processed item1" + assert results[1] == "Async processed item2" + assert results[2] == "Async processed item3" + + spans = instrumentation.get_finished_spans() + tool_span = next( + span for span in spans if span.attributes.get(SpanAttributes.AGENTOPS_SPAN_KIND) == SpanKind.TOOL + ) + assert tool_span.attributes.get(SpanAttributes.LLM_USAGE_TOOL_COST) == 0.04 + + def test_multiple_tool_calls(self, agent_class, instrumentation: InstrumentationTester): + """Test multiple calls to the same tool.""" + for i in range(3): + result = agent_class.process_item(f"item{i}") + assert result == f"Processed item{i}" + + spans = instrumentation.get_finished_spans() + tool_spans = [span for span in spans if span.attributes.get(SpanAttributes.AGENTOPS_SPAN_KIND) == SpanKind.TOOL] + assert len(tool_spans) == 3 + for span in tool_spans: + assert span.attributes.get(SpanAttributes.LLM_USAGE_TOOL_COST) == 0.01 + + @pytest.mark.asyncio + async def test_parallel_tool_calls(self, agent_class, instrumentation: InstrumentationTester): + """Test parallel execution of async tools.""" + results = await asyncio.gather( + agent_class.async_process_item("item1"), + agent_class.async_process_item("item2"), + agent_class.async_process_item("item3"), + ) + + assert len(results) == 3 + assert results[0] == "Async processed item1" + assert results[1] == "Async processed item2" + assert results[2] == "Async processed item3" + + spans = instrumentation.get_finished_spans() + tool_spans = [span for span in spans if span.attributes.get(SpanAttributes.AGENTOPS_SPAN_KIND) == SpanKind.TOOL] + assert len(tool_spans) == 3 + for span in tool_spans: + assert span.attributes.get(SpanAttributes.LLM_USAGE_TOOL_COST) == 0.02 + + def test_tool_without_cost(self, agent_class, instrumentation: InstrumentationTester): + """Test tool without cost parameter.""" + + @tool + def no_cost_tool(self): + return "No cost tool result" + + result = no_cost_tool(agent_class) + + assert result == "No cost tool result" + + spans = instrumentation.get_finished_spans() + tool_span = next( + span for span in spans if span.attributes.get(SpanAttributes.AGENTOPS_SPAN_KIND) == SpanKind.TOOL + ) + assert SpanAttributes.LLM_USAGE_TOOL_COST not in tool_span.attributes