From 50583092b710ec5b83d531f4eb070d1c4fe1c4fb Mon Sep 17 00:00:00 2001 From: fenilfaldu Date: Tue, 27 May 2025 19:53:53 +0530 Subject: [PATCH 1/6] Added the cost attribute in tool decorator along with the tests --- agentops/sdk/decorators/__init__.py | 4 +-- agentops/sdk/decorators/factory.py | 56 +++++++++++++++++++++++++---- agentops/semconv/span_attributes.py | 1 + docs/v1/concepts/sessions.mdx | 1 + docs/v2/concepts/decorators.mdx | 45 +++++++++++++++++++++++ docs/v2/usage/sdk-reference.mdx | 1 + 6 files changed, 99 insertions(+), 9 deletions(-) diff --git a/agentops/sdk/decorators/__init__.py b/agentops/sdk/decorators/__init__.py index 706bd4624..920b2fc2c 100644 --- a/agentops/sdk/decorators/__init__.py +++ b/agentops/sdk/decorators/__init__.py @@ -14,8 +14,8 @@ 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..221aa72d3 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, @@ -17,6 +19,7 @@ ) + def create_entity_decorator(entity_kind: str): """ Factory function that creates decorators for specific entity kinds. @@ -28,21 +31,24 @@ 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): # Create a proxy class that wraps the original class 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__() try: + _record_entity_input(self._agentops_active_span, args, kwargs) except Exception as e: logger.warning(f"Failed to record entity input: {e}") @@ -56,10 +62,12 @@ async def __aenter__(self): # 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 @@ -70,10 +78,12 @@ async def __aexit__(self, exc_type, exc_val, exc_tb): if hasattr(self, "_agentops_active_span") and hasattr(self, "_agentops_span_context_manager"): try: + _record_entity_output(self._agentops_active_span, self) except Exception as e: logger.warning(f"Failed to record entity output: {e}") + self._agentops_span_context_manager.__exit__(exc_type, exc_val, exc_tb) # Clear the span references after cleanup self._agentops_span_context_manager = None @@ -90,24 +100,35 @@ async def __aexit__(self, exc_type, exc_val, exc_tb): # Create the actual decorator wrapper function for functions @wrapt.decorator def wrapper(wrapped, instance, args, kwargs): + + # Skip instrumentation if tracer not initialized if not TracingCore.get_instance()._initialized: + return wrapped(*args, **kwargs) # Use provided name or function name operation_name = name or wrapped.__name__ + # Handle different types of functions (sync, async, generators) is_async = asyncio.iscoroutinefunction(wrapped) or inspect.iscoroutinefunction(wrapped) is_generator = inspect.isgeneratorfunction(wrapped) is_async_generator = inspect.isasyncgenfunction(wrapped) + + # 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 +137,15 @@ 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}") @@ -128,22 +154,30 @@ def wrapper(wrapped, instance, args, kwargs): # Handle async functions elif is_async: - 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}") try: + result = await wrapped(*args, **kwargs) try: + _record_entity_output(span, result) except Exception as e: 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 @@ -151,24 +185,32 @@ async def _wrapped_async(): # Handle sync functions else: + 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}") try: + result = wrapped(*args, **kwargs) try: + _record_entity_output(span, result) except Exception as e: 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 + raise # Return the wrapper for functions, we already returned WrappedClass for classes return wrapper(wrapped) # type: ignore diff --git a/agentops/semconv/span_attributes.py b/agentops/semconv/span_attributes.py index 79f0285a9..f0de6ff31 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.tool_cost" # Message attributes # see ./message.py for message-related attributes 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/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/usage/sdk-reference.mdx b/docs/v2/usage/sdk-reference.mdx index 7bd33a45b..6852bc7e9 100644 --- a/docs/v2/usage/sdk-reference.mdx +++ b/docs/v2/usage/sdk-reference.mdx @@ -145,6 +145,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. From 779ab4e2a481913b16406b23a5e810def2c59cd3 Mon Sep 17 00:00:00 2001 From: fenilfaldu Date: Tue, 27 May 2025 20:02:00 +0530 Subject: [PATCH 2/6] ruff check and tests added --- agentops/sdk/decorators/factory.py | 37 +------- agentops/semconv/span_attributes.py | 2 +- tests/unit/sdk/test_decorators.py | 139 +++++++++++++++++++++++++++- 3 files changed, 144 insertions(+), 34 deletions(-) diff --git a/agentops/sdk/decorators/factory.py b/agentops/sdk/decorators/factory.py index 221aa72d3..25e42532f 100644 --- a/agentops/sdk/decorators/factory.py +++ b/agentops/sdk/decorators/factory.py @@ -19,7 +19,6 @@ ) - def create_entity_decorator(entity_kind: str): """ Factory function that creates decorators for specific entity kinds. @@ -31,7 +30,7 @@ def create_entity_decorator(entity_kind: str): A decorator with optional arguments for name and version """ - def decorator(wrapped=None, *, name=None, version=None, cost=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, cost=cost) @@ -41,14 +40,12 @@ def decorator(wrapped=None, *, name=None, version=None, cost=None): # Create a proxy class that wraps the original class 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__() try: - _record_entity_input(self._agentops_active_span, args, kwargs) except Exception as e: logger.warning(f"Failed to record entity input: {e}") @@ -62,12 +59,11 @@ async def __aenter__(self): # 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 @@ -78,12 +74,10 @@ async def __aexit__(self, exc_type, exc_val, exc_tb): if hasattr(self, "_agentops_active_span") and hasattr(self, "_agentops_span_context_manager"): try: - _record_entity_output(self._agentops_active_span, self) except Exception as e: logger.warning(f"Failed to record entity output: {e}") - self._agentops_span_context_manager.__exit__(exc_type, exc_val, exc_tb) # Clear the span references after cleanup self._agentops_span_context_manager = None @@ -100,34 +94,25 @@ async def __aexit__(self, exc_type, exc_val, exc_tb): # Create the actual decorator wrapper function for functions @wrapt.decorator def wrapper(wrapped, instance, args, kwargs): - - # Skip instrumentation if tracer not initialized if not TracingCore.get_instance()._initialized: - return wrapped(*args, **kwargs) # Use provided name or function name operation_name = name or wrapped.__name__ - # Handle different types of functions (sync, async, generators) is_async = asyncio.iscoroutinefunction(wrapped) or inspect.iscoroutinefunction(wrapped) is_generator = inspect.isgeneratorfunction(wrapped) is_async_generator = inspect.isasyncgenfunction(wrapped) - - # Handle generator functions if is_generator: - 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}") @@ -137,14 +122,11 @@ def wrapper(wrapped, instance, args, kwargs): # Handle async generator functions elif is_async_generator: - 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}") @@ -154,24 +136,20 @@ def wrapper(wrapped, instance, args, kwargs): # Handle async functions elif is_async: + 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}") try: - result = await wrapped(*args, **kwargs) try: - _record_entity_output(span, result) except Exception as e: logger.warning(f"Failed to record entity output: {e}") @@ -185,24 +163,19 @@ async def _wrapped_async(): # Handle sync functions else: - 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}") try: - result = wrapped(*args, **kwargs) try: - _record_entity_output(span, result) except Exception as e: logger.warning(f"Failed to record entity output: {e}") @@ -210,7 +183,7 @@ async def _wrapped_async(): except Exception as e: logger.error(f"Error in sync function execution: {e}") span.record_exception(e) - raise + raise # Return the wrapper for functions, we already returned WrappedClass for classes return wrapper(wrapped) # type: ignore diff --git a/agentops/semconv/span_attributes.py b/agentops/semconv/span_attributes.py index f0de6ff31..231ea2cf7 100644 --- a/agentops/semconv/span_attributes.py +++ b/agentops/semconv/span_attributes.py @@ -65,7 +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.tool_cost" + LLM_USAGE_TOOL_COST = "gen_ai.usage.tool_cost" # Message attributes # see ./message.py for message-related attributes 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 From 571a37cebec427a2c7de471088f9e8631b683f9b Mon Sep 17 00:00:00 2001 From: fenilfaldu Date: Tue, 27 May 2025 20:08:04 +0530 Subject: [PATCH 3/6] code cleanup --- docs/v1/concepts/sessions.mdx | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/docs/v1/concepts/sessions.mdx b/docs/v1/concepts/sessions.mdx index d9996834a..60a503c8f 100644 --- a/docs/v1/concepts/sessions.mdx +++ b/docs/v1/concepts/sessions.mdx @@ -27,7 +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` From 3d30c27b36149afb60864735706083c384b05e69 Mon Sep 17 00:00:00 2001 From: fenilfaldu Date: Tue, 27 May 2025 20:19:08 +0530 Subject: [PATCH 4/6] minor changes --- agentops/sdk/decorators/__init__.py | 1 + agentops/sdk/decorators/factory.py | 6 ------ agentops/semconv/span_attributes.py | 2 +- 3 files changed, 2 insertions(+), 7 deletions(-) diff --git a/agentops/sdk/decorators/__init__.py b/agentops/sdk/decorators/__init__.py index 920b2fc2c..c17b08aa7 100644 --- a/agentops/sdk/decorators/__init__.py +++ b/agentops/sdk/decorators/__init__.py @@ -16,6 +16,7 @@ session = create_entity_decorator(SpanKind.SESSION) tool = create_entity_decorator(SpanKind.TOOL) operation = task + __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 25e42532f..582005ccd 100644 --- a/agentops/sdk/decorators/factory.py +++ b/agentops/sdk/decorators/factory.py @@ -54,9 +54,6 @@ 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 @@ -69,9 +66,6 @@ async def __aenter__(self): 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) diff --git a/agentops/semconv/span_attributes.py b/agentops/semconv/span_attributes.py index 231ea2cf7..bb5849888 100644 --- a/agentops/semconv/span_attributes.py +++ b/agentops/semconv/span_attributes.py @@ -65,7 +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.tool_cost" + LLM_USAGE_TOOL_COST = "gen_ai.usage.total_cost" # Message attributes # see ./message.py for message-related attributes From 22a669380317040b9cff30d8d084aa5fc46b1499 Mon Sep 17 00:00:00 2001 From: fenilfaldu Date: Tue, 27 May 2025 20:43:40 +0530 Subject: [PATCH 5/6] updated the testcase --- tests/unit/sdk/test_decorators.py | 9 +++++---- 1 file changed, 5 insertions(+), 4 deletions(-) diff --git a/tests/unit/sdk/test_decorators.py b/tests/unit/sdk/test_decorators.py index 80e2e7ddf..6408b67c8 100644 --- a/tests/unit/sdk/test_decorators.py +++ b/tests/unit/sdk/test_decorators.py @@ -708,10 +708,11 @@ async def test_async_generator_tool_cost(self, agent_class, instrumentation: Ins 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 + tool_span = [span for span in spans if span.attributes.get(SpanAttributes.AGENTOPS_SPAN_KIND) == SpanKind.TOOL] + assert len(tool_span) == 4 # Only one span for the generator + assert tool_span[0].attributes.get(SpanAttributes.LLM_USAGE_TOOL_COST) == 0.02 + assert tool_span[3].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.""" From 5be9f735dbbcc81698cb5c69dc1fa90b360d1292 Mon Sep 17 00:00:00 2001 From: fenilfaldu Date: Tue, 27 May 2025 20:46:17 +0530 Subject: [PATCH 6/6] ruff checks --- tests/unit/sdk/test_decorators.py | 1 - 1 file changed, 1 deletion(-) diff --git a/tests/unit/sdk/test_decorators.py b/tests/unit/sdk/test_decorators.py index 6408b67c8..96824d0fe 100644 --- a/tests/unit/sdk/test_decorators.py +++ b/tests/unit/sdk/test_decorators.py @@ -713,7 +713,6 @@ async def test_async_generator_tool_cost(self, agent_class, instrumentation: Ins assert tool_span[0].attributes.get(SpanAttributes.LLM_USAGE_TOOL_COST) == 0.02 assert tool_span[3].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):