Add span functions

README.md

@@ -102,9 +102,13 @@ async def main(input: dict, context: dict):
     return result["output"]
 ```
 
-### Using Custom Decorators (Any Framework)
+### Manual Trace and Span Capture
 
-For frameworks beyond LangGraph, use trace decorators to capture custom spans:
+For frameworks beyond LangGraph, you can manually capture traces using either **decorators** (wrap functions) or **functions** (record spans after execution).
+
+#### Using Decorators
+
+Use `@trace_llm`, `@trace_tool`, and `@trace_retriever` decorators to automatically capture spans when functions are called:
 
 ```python
 from gradient_adk import entrypoint, trace_llm, trace_tool, trace_retriever
@@ -117,7 +121,7 @@ async def search_knowledge_base(query: str):
 
 @trace_llm("generate_response")
 async def generate_response(prompt: str):
-    # LLM spans capture model calls with token usage
+    # LLM spans capture model calls
     response = await llm.generate(prompt)
     return response
 
@@ -134,6 +138,36 @@ async def main(input: dict, context: dict):
     return response
 ```
 
+#### Using Functions
+
+Use `add_llm_span`, `add_tool_span`, and `add_retriever_span` functions to manually record spans after execution. This is useful when you can't wrap a function with a decorator or need more control over what gets recorded:
+
+```python
+from gradient_adk import entrypoint, add_llm_span, add_tool_span, add_retriever_span
+
+@entrypoint
+async def main(input: dict, context: dict):
+    # Perform retrieval and record the span
+    results = await vector_db.search(input["query"])
+    add_retriever_span("vector_search", inputs={"query": input["query"]}, output=results)
+
+    # Perform calculation and record the span
+    result = 5 + 10
+    add_tool_span("calculate", inputs={"x": 5, "y": 10}, output=result)
+
+    # Call LLM and record with additional metadata
+    response = await llm.generate(f"Context: {results}")
+    add_llm_span(
+        "generate_response",
+        inputs={"prompt": f"Context: {results}"},
+        output=response,
+        model_name="gpt-4",      # Optional: model name
+        ttft_ms=150.5,           # Optional: time to first token
+    )
+
+    return response
+```
+
 ### Streaming Responses
 
 The runtime supports streaming responses with automatic trace capture:
@@ -204,7 +238,9 @@ The ADK runtime automatically captures detailed traces:
 - **Errors**: Full exception details and stack traces
 - **Streaming Responses**: Individual chunks and aggregated outputs
 
-### Available Decorators
+### Available Decorators and Functions
+
+**Decorators** - Wrap functions to automatically capture spans:
 ```python
 from gradient_adk import trace_llm, trace_tool, trace_retriever
 
@@ -213,7 +249,17 @@ from gradient_adk import trace_llm, trace_tool, trace_retriever
 @trace_retriever("db_search")  # For retrieval/search operations
 ```
 
-These decorators are used to log steps or spans of your agent workflow that are not automatically captured.  These will log things like the input, output, and step duration and make them available in your agent's traces and for use in agent evaluations.
+**Functions** - Manually record spans after execution:
+```python
+from gradient_adk import add_llm_span, add_tool_span, add_retriever_span
+
+# Record spans with name, inputs, and output
+add_llm_span("model_call", inputs={"prompt": "..."}, output="response", model_name="gpt-4")
+add_tool_span("calculator", inputs={"x": 5, "y": 10}, output=15)
+add_retriever_span("db_search", inputs={"query": "..."}, output=[...])
+```
+
+These are used to log steps or spans of your agent workflow that are not automatically captured. They will log things like the input, output, and step duration and make them available in your agent's traces and for use in agent evaluations.
 
 ### Viewing Traces
 Traces are:
@@ -234,20 +280,6 @@ export GRADIENT_MODEL_ACCESS_KEY=your_gradient_key
 export GRADIENT_VERBOSE=1
 ```
 
-## Project Structure
-
-```
-my-agent/
-├── main.py              # Agent entrypoint with @entrypoint decorator
-├── .gradient/agent.yml  # Agent configuration (auto-generated)
-├── requirements.txt     # Python dependencies
-├── .env                 # Environment variables (not committed)
-├── agents/              # Agent implementations
-│   └── my_agent.py
-└── tools/               # Custom tools
-    └── my_tool.py
-```
-
 ## Framework Compatibility
 
 The Gradient ADK is designed to work with any Python-based AI agent framework:
@@ -267,4 +299,4 @@ The Gradient ADK is designed to work with any Python-based AI agent framework:
 
 ## License
 
-Licensed under the Apache License 2.0. See [LICENSE](./LICENSE)
+Licensed under the Apache License 2.0. See [LICENSE](./LICENSE)

gradient_adk/__init__.py

@@ -4,17 +4,25 @@
 """
 
 from .decorator import entrypoint
-from .tracing import (  # manual tracing decorators
+from .tracing import (  # manual tracing decorators and functions
     trace_llm,
     trace_retriever,
     trace_tool,
+    add_llm_span,
+    add_retriever_span,
+    add_tool_span,
 )
 
 __all__ = [
     "entrypoint",
+    # Decorators
     "trace_llm",
     "trace_retriever",
     "trace_tool",
+    # Functions
+    "add_llm_span",
+    "add_retriever_span",
+    "add_tool_span",
 ]
 
 __version__ = "0.0.5"

gradient_adk/tracing.py

@@ -1,9 +1,10 @@
-"""Tracing decorators for manual span tracking.
+"""Tracing decorators and functions for manual span tracking.
 
-These decorators allow developers to instrument their custom agent functions
-with the same kind of tracing automatically provided for some other frameworks.
+These decorators and functions allow developers to instrument their custom agent
+functions with the same kind of tracing automatically provided for some other
+frameworks.
 
-Example usage:
+Decorator-based usage:
     from gradient_adk import entrypoint, trace_llm, trace_tool, trace_retriever
 
     @trace_retriever("fetch_data")
@@ -27,6 +28,28 @@ async def my_agent(input: dict, context: dict):
         result = await calculate(5, 10)
         response = await call_model(data["prompt"])
         return {"response": response}
+
+Function-based usage (for manual span creation):
+    from gradient_adk import entrypoint, add_llm_span, add_tool_span, add_retriever_span
+
+    @entrypoint
+    async def my_agent(input: dict, context: dict):
+        # Do work and then record spans manually
+        query = input["query"]
+
+        # Record a retriever span
+        results = await fetch_documents(query)
+        add_retriever_span("fetch_docs", inputs={"query": query}, output=results)
+
+        # Record a tool span
+        calculation = calculate(5, 10)
+        add_tool_span("calculate", inputs={"x": 5, "y": 10}, output=calculation)
+
+        # Record an LLM span
+        response = await call_llm(results)
+        add_llm_span("generate_response", inputs={"context": results}, output=response)
+
+        return {"response": response}
 """
 
 from __future__ import annotations
@@ -443,3 +466,160 @@ async def search(query: str) -> list:
             return results
     """
     return _trace_base(name, span_type=SpanType.TOOL)
+
+
+def _add_span(
+    name: str,
+    inputs: Any,
+    output: Any,
+    span_type: SpanType,
+    extra_metadata: Optional[Dict[str, Any]] = None,
+) -> None:
+    """
+    Internal helper to add a completed span to the tracker.
+
+    Args:
+        name: Name for the span.
+        inputs: The inputs to record for this span.
+        output: The output to record for this span.
+        span_type: Type of span (LLM, TOOL, or RETRIEVER).
+        extra_metadata: Additional metadata fields to attach to the span.
+    """
+    tracker = get_tracker()
+    if not tracker:
+        # No tracker available, silently skip
+        return
+
+    # Snapshot inputs and output
+    inputs_snapshot = _freeze(inputs)
+    output_snapshot = _freeze(output)
+
+    # Create span
+    span = _create_span(name, inputs_snapshot)
+
+    # Mark span type
+    meta = _ensure_meta(span)
+    if span_type == SpanType.LLM:
+        meta["is_llm_call"] = True
+    elif span_type == SpanType.TOOL:
+        meta["is_tool_call"] = True
+    elif span_type == SpanType.RETRIEVER:
+        meta["is_retriever_call"] = True
+
+    # Add any extra metadata
+    if extra_metadata:
+        for key, value in extra_metadata.items():
+            meta[key] = value
+
+    # Record start and end
+    tracker.on_node_start(span)
+    tracker.on_node_end(span, output_snapshot)
+
+
+def add_llm_span(
+    name: str,
+    inputs: Any,
+    output: Any,
+    *,
+    model_name: Optional[str] = None,
+    ttft_ms: Optional[float] = None,
+    **extra_metadata: Any,
+) -> None:
+    """
+    Add an LLM call span with the given name, inputs, and output.
+
+    Use this function to manually record an LLM call span.
+
+    Args:
+        name: Name for the span (e.g., "openai_call", "generate_response").
+        inputs: The inputs to the LLM call (e.g., prompt, messages).
+        output: The output from the LLM call (e.g., response text).
+        model_name: Optional name of the LLM model used.
+        ttft_ms: Optional time to first token in milliseconds.
+        **extra_metadata: Any additional metadata fields to attach to the span.
+
+    Example:
+        response = await call_llm(prompt)
+        add_llm_span(
+            "generate_response",
+            inputs={"prompt": prompt},
+            output=response,
+            model_name="gpt-4",
+            ttft_ms=150.5,
+            temperature=0.7,
+        )
+    """
+    metadata: Dict[str, Any] = {}
+    if model_name is not None:
+        metadata["model_name"] = model_name
+    if ttft_ms is not None:
+        metadata["ttft_ms"] = ttft_ms
+    metadata.update(extra_metadata)
+
+    _add_span(name, inputs, output, SpanType.LLM, metadata if metadata else None)
+
+
+def add_retriever_span(
+    name: str,
+    inputs: Any,
+    output: Any,
+    **extra_metadata: Any,
+) -> None:
+    """
+    Add a retriever call span with the given name, inputs, and output.
+
+    Use this function to manually record a retriever span without using a decorator.
+
+    Args:
+        name: Name for the span (e.g., "vector_search", "fetch_docs").
+        inputs: The inputs to the retriever (e.g., query, filters).
+        output: The output from the retriever (e.g., list of documents).
+        **extra_metadata: Any additional metadata fields to attach to the span.
+
+    Example:
+        results = await fetch_documents(query)
+        add_retriever_span(
+            "fetch_docs",
+            inputs={"query": query},
+            output=results,
+            num_results=len(results),
+        )
+    """
+    _add_span(
+        name,
+        inputs,
+        output,
+        SpanType.RETRIEVER,
+        extra_metadata if extra_metadata else None,
+    )
+
+
+def add_tool_span(
+    name: str,
+    inputs: Any,
+    output: Any,
+    **extra_metadata: Any,
+) -> None:
+    """
+    Add a tool call span with the given name, inputs, and output.
+
+    Use this function to manually record a tool call span without using a decorator.
+
+    Args:
+        name: Name for the span (e.g., "calculate", "search_database").
+        inputs: The inputs to the tool (e.g., parameters).
+        output: The output from the tool (e.g., result).
+        **extra_metadata: Any additional metadata fields to attach to the span.
+
+    Example:
+        result = calculate(5, 10)
+        add_tool_span(
+            "calculate",
+            inputs={"x": 5, "y": 10},
+            output=result,
+            execution_time_ms=12.5,
+        )
+    """
+    _add_span(
+        name, inputs, output, SpanType.TOOL, extra_metadata if extra_metadata else None
+    )