Enhance documentation and functionality for trace management

docs/mint.json

@@ -188,7 +188,9 @@
         "v2/usage/advanced-configuration",
         "v2/usage/tracking-llm-calls",
         "v2/usage/tracking-agents",
-        "v2/usage/recording-operations"
+        "v2/usage/recording-operations",
+        "v2/usage/trace-decorator",
+        "v2/usage/manual-trace-control"
       ],
       "version": "v2"
     },

docs/v2/integrations/google_adk.mdx

@@ -21,73 +21,162 @@ AgentOps provides seamless integration with [Google Agent Development Kit (ADK)]
       ```
     </CodeGroup>
 
-## Basic Usage
+## Usage
 
 Initialize AgentOps at the beginning of your application to automatically track all Google ADK agent interactions:
 
     <CodeGroup>
-```python Basic Usage
-import asyncio
-import uuid
-import os
-from google.genai import types
+```python Usage
+# --- Full example code demonstrating LlmAgent with Tools vs. Output Schema ---
+import json # Needed for pretty printing dicts
 
+from google.adk.agents import LlmAgent
+from google.adk.runners import Runner
+from google.adk.sessions import InMemorySessionService
+from google.genai import types
+from pydantic import BaseModel, Field
+import asyncio
 import agentops
-from google.adk.agents import Agent
-from google.adk.runners import InMemoryRunner
-from google.adk.agents.run_config import RunConfig, StreamingMode
-
-# Initialize AgentOps
-agentops.init(api_key="<INSERT YOUR API KEY HERE>")
-
-# Create a simple agent with no tools
-agent = Agent(
-    name="simple_agent",
-    model="gemini-1.5-flash",
-    instruction="You are a helpful assistant that provides clear and concise answers.",
+
+agentops.init("your-api-key")
+
+# --- 1. Define Constants ---
+APP_NAME = "agent_comparison_app"
+USER_ID = "test_user_456"
+SESSION_ID_TOOL_AGENT = "session_tool_agent_xyz"
+SESSION_ID_SCHEMA_AGENT = "session_schema_agent_xyz"
+MODEL_NAME = "gemini-2.0-flash"
+
+# --- 2. Define Schemas ---
+
+# Input schema used by both agents
+class CountryInput(BaseModel):
+    country: str = Field(description="The country to get information about.")
+
+# Output schema ONLY for the second agent
+class CapitalInfoOutput(BaseModel):
+    capital: str = Field(description="The capital city of the country.")
+    # Note: Population is illustrative; the LLM will infer or estimate this
+    # as it cannot use tools when output_schema is set.
+    population_estimate: str = Field(description="An estimated population of the capital city.")
+
+# --- 3. Define the Tool (Only for the first agent) ---
+def get_capital_city(country: str) -> str:
+    """Retrieves the capital city of a given country."""
+    print(f"\n-- Tool Call: get_capital_city(country='{country}') --")
+    country_capitals = {
+        "united states": "Washington, D.C.",
+        "canada": "Ottawa",
+        "france": "Paris",
+        "japan": "Tokyo",
+    }
+    result = country_capitals.get(country.lower(), f"Sorry, I couldn't find the capital for {country}.")
+    print(f"-- Tool Result: '{result}' --")
+    return result
+
+# --- 4. Configure Agents ---
+
+# Agent 1: Uses a tool and output_key
+capital_agent_with_tool = LlmAgent(
+    model=MODEL_NAME,
+    name="capital_agent_tool",
+    description="Retrieves the capital city using a specific tool.",
+    instruction="""You are a helpful agent that provides the capital city of a country using a tool.
+The user will provide the country name in a JSON format like {"country": "country_name"}.
+1. Extract the country name.
+2. Use the `get_capital_city` tool to find the capital.
+3. Respond clearly to the user, stating the capital city found by the tool.
+""",
+    tools=[get_capital_city],
+    input_schema=CountryInput,
+    output_key="capital_tool_result", # Store final text response
 )
 
-# Create a runner
-runner = InMemoryRunner(
-    agent=agent,
-    app_name="simple-example",
+# Agent 2: Uses output_schema (NO tools possible)
+structured_info_agent_schema = LlmAgent(
+    model=MODEL_NAME,
+    name="structured_info_agent_schema",
+    description="Provides capital and estimated population in a specific JSON format.",
+    instruction=f"""You are an agent that provides country information.
+The user will provide the country name in a JSON format like {{"country": "country_name"}}.
+Respond ONLY with a JSON object matching this exact schema:
+{json.dumps(CapitalInfoOutput.model_json_schema(), indent=2)}
+Use your knowledge to determine the capital and estimate the population. Do not use any tools.
+""",
+    # *** NO tools parameter here - using output_schema prevents tool use ***
+    input_schema=CountryInput,
+    output_schema=CapitalInfoOutput, # Enforce JSON output structure
+    output_key="structured_info_result", # Store final JSON response
 )
 
-# Setup session
-user_id = f"user-{uuid.uuid4().hex[:8]}"
-session_id = f"session-{uuid.uuid4().hex[:8]}"
-runner.session_service.create_session(
-    app_name="simple-example",
-    user_id=user_id,
-    session_id=session_id,
+# --- 5. Set up Session Management and Runners ---
+session_service = InMemorySessionService()
+
+# Create a runner for EACH agent
+capital_runner = Runner(
+    agent=capital_agent_with_tool,
+    app_name=APP_NAME,
+    session_service=session_service
+)
+structured_runner = Runner(
+    agent=structured_info_agent_schema,
+    app_name=APP_NAME,
+    session_service=session_service
 )
 
-# Run the agent with a user message
-async def run_agent():
-    message = "What are three benefits of artificial intelligence?"
-
-    content = types.Content(
-        role="user",
-        parts=[types.Part(text=message)],
-    )
-
-    run_config = RunConfig(
-        streaming_mode=StreamingMode.NONE,
-    )
+# --- 6. Define Agent Interaction Logic ---
+async def call_agent_and_print(
+    runner_instance: Runner,
+    agent_instance: LlmAgent,
+    session_id: str,
+    query_json: str
+):
+    """Sends a query to the specified agent/runner and prints results."""
+    print(f"\n>>> Calling Agent: '{agent_instance.name}' | Query: {query_json}")
+
+    user_content = types.Content(role='user', parts=[types.Part(text=query_json)])
+
+    final_response_content = "No final response received."
+    async for event in runner_instance.run_async(user_id=USER_ID, session_id=session_id, new_message=user_content):
+        # print(f"Event: {event.type}, Author: {event.author}") # Uncomment for detailed logging
+        if event.is_final_response() and event.content and event.content.parts:
+            # For output_schema, the content is the JSON string itself
+            final_response_content = event.content.parts[0].text
+
+    print(f"<<< Agent '{agent_instance.name}' Response: {final_response_content}")
+
+    current_session = await session_service.get_session(app_name=APP_NAME,
+                                                         user_id=USER_ID,
+                                                         session_id=session_id)
+    stored_output = current_session.state.get(agent_instance.output_key)
+
+    # Pretty print if the stored output looks like JSON (likely from output_schema)
+    print(f"--- Session State ['{agent_instance.output_key}']: ", end="")
+    try:
+        # Attempt to parse and pretty print if it's JSON
+        parsed_output = json.loads(stored_output)
+        print(json.dumps(parsed_output, indent=2))
+    except (json.JSONDecodeError, TypeError):
+         # Otherwise, print as string
+        print(stored_output)
+    print("-" * 30)
+
+
+# --- 7. Run Interactions ---
+async def main():
+    # Create sessions
+    await session_service.create_session(app_name=APP_NAME, user_id=USER_ID, session_id=SESSION_ID_TOOL_AGENT)
+    await session_service.create_session(app_name=APP_NAME, user_id=USER_ID, session_id=SESSION_ID_SCHEMA_AGENT)
 
-    async for event in runner.run_async(
-        user_id=user_id,
-        session_id=session_id,
-        new_message=content,
-        run_config=run_config,
-    ):
-        if hasattr(event, 'content') and event.content and event.content.parts:
-            for part in event.content.parts:
-                if hasattr(part, 'text') and part.text:
-                    print(part.text)
-
-# Run the agent
-asyncio.run(run_agent())
+    print("--- Testing Agent with Tool ---")
+    await call_agent_and_print(capital_runner, capital_agent_with_tool, SESSION_ID_TOOL_AGENT, '{"country": "France"}')
+    await call_agent_and_print(capital_runner, capital_agent_with_tool, SESSION_ID_TOOL_AGENT, '{"country": "Canada"}')
+
+    print("\n\n--- Testing Agent with Output Schema (No Tool Use) ---")
+    await call_agent_and_print(structured_runner, structured_info_agent_schema, SESSION_ID_SCHEMA_AGENT, '{"country": "France"}')
+    await call_agent_and_print(structured_runner, structured_info_agent_schema, SESSION_ID_SCHEMA_AGENT, '{"country": "Japan"}')
+
+asyncio.run(main())
       ```
     </CodeGroup>
 

docs/v2/introduction.mdx

@@ -35,22 +35,32 @@ Observability and monitoring for your AI agents and LLM apps. And we do it all i
 </CodeGroup>
 ... that logs everything back to your AgentOps Dashboard.
 
-That's it! AgentOps will automatically instrument your code and start tracking sessions.
+That's it! AgentOps will automatically instrument your code and start tracking traces.
 
-Need more control? You can disable automatic session creation and manage sessions explicitly:
+Need more control? You can disable automatic session creation and manage traces manually:
 
 <CodeGroup>
 	```python python
 	import agentops
 	agentops.init(<INSERT YOUR API KEY HERE>, auto_start_session=False)
 
-	# Later, when you're ready to start a session:
-	agentops.start_session("my-workflow-session")
+	# Later, when you're ready to start a trace:
+	trace = agentops.start_trace("my-workflow-trace")
 
 	# Your code here
 	# ...
 
-	# Sessions automatically end when your program exits
+	# End the trace when done
+	agentops.end_trace(trace, "Success")
+	```
+</CodeGroup>
+
+You can also set a custom trace name during initialization:
+
+<CodeGroup>
+	```python python
+	import agentops
+	agentops.init(<INSERT YOUR API KEY HERE>, trace_name="custom-trace-name")
 	```
 </CodeGroup>
 

docs/v2/quickstart.mdx

@@ -25,7 +25,7 @@ import EnvTooltip from '/snippets/add-env-tooltip.mdx'
         agentops.init(<INSERT YOUR API KEY HERE>)
         ```
       </CodeGroup>
-      That's it! These two lines automatically instrument your code and start tracking sessions. Sessions automatically end when your program exits.
+      That's it! These two lines automatically instrument your code and start tracking traces. Traces automatically end when your program exits.
     <EnvTooltip />
   </Step>
   <Step title="Run your agent">
@@ -84,21 +84,43 @@ import EnvTooltip from '/snippets/add-env-tooltip.mdx'
   ```
 
 </Card>
-<Card icon="stop" title="Creating Sessions">
-  Create a session to group all your agent operations by using the `@session` decorator. Sessions serve as the root span for all operations.
+<Card icon="wrench" title="Track Tools">
+  Track tool usage and costs with the `@tool` decorator. You can specify costs to get total cost tracking directly in your dashboard summary.
   ```python python
-  # Create a session
-  from agentops.sdk.decorators import session
+  # Track tool usage with cost
+  from agentops.sdk.decorators import tool
 
-  @session
+  @tool(cost=0.05)
+  def web_search(query):
+      # Tool logic here
+      return f"Search results for: {query}"
+
+  @tool
+  def calculate(expression):
+      # Tool without cost tracking
+      return eval(expression)
+  ```
+
+</Card>
+<Card icon="play" title="Creating Traces">
+  Create custom traces to group operations using the `@trace` decorator, or manage traces manually for more control.
+  ```python python
+  # Create a trace with decorator
+  from agentops.sdk.decorators import trace
+
+  @trace
   def my_workflow():
-      # Your session code here
+      # Your workflow code here
       agent = MyAgent("research-agent")
       result = agent.perform_task("data analysis")
       return result
 
-  # Run the session
-  my_workflow()
+  # Or manage traces manually
+  import agentops
+
+  trace = agentops.start_trace("custom-trace")
+  # Your code here
+  agentops.end_trace(trace, "Success")
   ```
 </Card>
 </CardGroup>
@@ -110,10 +132,15 @@ Here is the complete code from the sections above
 
 ```python python
 import agentops
-from agentops.sdk.decorators import agent, operation
+from agentops.sdk.decorators import agent, operation, tool, trace
 
 # Initialize AgentOps
-agentops.init(<INSERT YOUR API KEY HERE>)
+agentops.init(<INSERT YOUR API KEY HERE>, auto_start_session=False)
+
+# Create a tool with cost tracking
+@tool(cost=0.05)
+def web_search(query):
+    return f"Search results for: {query}"
 
 # Create an agent class
 @agent
@@ -123,9 +150,11 @@ class MyAgent:
 
     @operation
     def perform_task(self, task):
-        # Agent task logic here
-        return f"Completed {task}"
+        # Use a tool within the agent
+        search_results = web_search(f"research {task}")
+        return f"Completed {task} with results: {search_results}"
 
+@trace
 def run_agent_task(task_name):
     agent = MyAgent("research-agent")
     result = agent.perform_task(task_name)