Nicer agentic example

Author: piotrekno1

examples/agentic_example.py

@@ -14,10 +14,13 @@
   - Tracks LLM inference automatically via the OpenAI integration
 
 Run with: uv run agentic_example.py
-Requires: OPENAI_API_KEY environment variable. Set WILDEDGE_DSN to send events.
+Requires: OPENROUTER_API_KEY environment variable. Set WILDEDGE_DSN to send events.
 """
 
 import json
+import os
+import time
+import uuid
 
 from openai import OpenAI
 
@@ -28,7 +31,10 @@
     integrations="openai",
 )
 
-openai_client = OpenAI()
+openai_client = OpenAI(
+    base_url="https://openrouter.ai/api/v1",
+    api_key=os.getenv("OPENROUTER_API_KEY"),
+)
 
 # --- Tools -------------------------------------------------------------------
 
@@ -65,11 +71,14 @@
 
 
 def get_weather(city: str) -> str:
-    # Stub: replace with a real weather API call.
+    # ~150ms to simulate a real weather API call.
+    time.sleep(0.15)
     return json.dumps({"city": city, "temperature_c": 18, "condition": "partly cloudy"})
 
 
 def calculator(expression: str) -> str:
+    # ~60ms to simulate a remote computation call.
+    time.sleep(0.06)
     try:
         result = eval(expression, {"__builtins__": {}})  # noqa: S307
         return json.dumps({"expression": expression, "result": result})
@@ -97,8 +106,23 @@ def call_tool(name: str, arguments: dict) -> str:
     return result
 
 
+def retrieve_context(query: str) -> str:
+    """Fetch relevant context from the vector store (~120ms)."""
+    with we.span(
+        kind="retrieval",
+        name="vector_search",
+        input_summary=query[:200],
+    ) as span:
+        time.sleep(0.12)
+        result = f"[context: background knowledge relevant to '{query[:40]}']"
+        span.output_summary = result
+    return result
+
+
 def run_agent(task: str, step_index: int, messages: list) -> str:
-    messages.append({"role": "user", "content": task})
+    # Fetch context before the first reasoning step, include it in the user turn.
+    context = retrieve_context(task)
+    messages.append({"role": "user", "content": f"{task}\n\nContext: {context}"})
 
     while True:
         with we.span(
@@ -108,15 +132,16 @@ def run_agent(task: str, step_index: int, messages: list) -> str:
             input_summary=task[:200],
         ) as span:
             response = openai_client.chat.completions.create(
-                model="gpt-4o",
+                model="qwen/qwen3.5-flash-02-23",
                 messages=messages,
                 tools=TOOLS,
                 tool_choice="auto",
+                max_tokens=512,
             )
             choice = response.choices[0]
             span.output_summary = choice.finish_reason
 
-        messages.append(choice.message)
+        messages.append(choice.message.model_dump(exclude_none=True))
 
         if choice.finish_reason == "tool_calls":
             step_index += 1
@@ -130,8 +155,11 @@ def run_agent(task: str, step_index: int, messages: list) -> str:
                         "content": result,
                     }
                 )
+                # Not instrumented: context window update between tool calls (~80ms).
+                # Shows up as a gap stripe in the trace view.
+                time.sleep(0.08)
         else:
-            return choice.message.content
+            return choice.message.content or ""
 
 
 # --- Main --------------------------------------------------------------------
@@ -144,11 +172,10 @@ def run_agent(task: str, step_index: int, messages: list) -> str:
 system_prompt = "You are a helpful assistant. Use tools when needed."
 messages = [{"role": "system", "content": system_prompt}]
 
-with we.trace(agent_id="demo-agent", run_id="example-run-001"):
+with we.trace(agent_id="demo-agent", run_id=str(uuid.uuid4())):
     for i, task in enumerate(TASKS, start=1):
         print(f"\nTask {i}: {task}")
         reply = run_agent(task, step_index=i, messages=messages)
         print(f"Reply: {reply}")
 
 we.flush()
-print("\nDone. Events flushed to WildEdge.")

tests/test_tracing.py

@@ -1,7 +1,8 @@
 from __future__ import annotations
 
+from wildedge.client import SpanContextManager
 from wildedge.model import ModelHandle, ModelInfo
-from wildedge.tracing import span_context, trace_context
+from wildedge.tracing import get_span_context, span_context, trace_context
 
 
 def test_track_inference_uses_trace_context():
@@ -36,3 +37,57 @@ def publish(event: dict) -> None:
     assert events[0]["agent_id"] == "agent-1"
     assert events[0]["step_index"] == 2
     assert events[0]["attributes"] == {"trace_key": "trace_val", "span_key": 2}
+
+
+class _FakeClient:
+    def __init__(self, events: list[dict]) -> None:
+        self._events = events
+
+    def track_span(self, **kwargs) -> str:
+        self._events.append(kwargs)
+        return kwargs.get("span_id", "")
+
+
+def test_span_root_has_no_parent():
+    """A root span must not reference itself as its own parent."""
+    events: list[dict] = []
+    client = _FakeClient(events)
+
+    with SpanContextManager(client, kind="agent_step", name="root"):
+        pass
+
+    assert len(events) == 1
+    assert events[0]["parent_span_id"] is None
+
+
+def test_span_context_restored_after_exit():
+    """The active span context must revert to the parent after a span exits."""
+    events: list[dict] = []
+    client = _FakeClient(events)
+
+    with span_context(span_id="parent-span"):
+        with SpanContextManager(client, kind="agent_step", name="child"):
+            inner_id = get_span_context().span_id
+
+        assert get_span_context().span_id == "parent-span"
+
+    assert inner_id != "parent-span"
+    assert events[0]["parent_span_id"] == "parent-span"
+    assert events[0]["span_id"] != "parent-span"
+
+
+def test_nested_spans_correct_parent_chain():
+    """Nested spans must each point to their direct parent, not themselves."""
+    events: list[dict] = []
+    client = _FakeClient(events)
+
+    with SpanContextManager(client, kind="agent_step", name="outer") as outer:
+        with SpanContextManager(client, kind="tool", name="inner") as inner:
+            pass
+
+    assert len(events) == 2
+    inner_ev, outer_ev = events[0], events[1]
+    assert inner_ev["span_id"] == inner.span_id
+    assert inner_ev["parent_span_id"] == outer.span_id
+    assert outer_ev["span_id"] == outer.span_id
+    assert outer_ev["parent_span_id"] is None

wildedge/client.py

@@ -175,6 +175,12 @@ def __exit__(self, exc_type, exc_val, exc_tb):
             return False
         duration_ms = elapsed_ms(self._t0)
         status = "error" if exc_type else self.status
+        # Restore parent span context before emitting, so _merge_correlation_fields
+        # sees the parent context rather than this span (which would make the span
+        # appear as its own parent).
+        if self._span_token is not None:
+            _reset_span_context(self._span_token)
+            self._span_token = None
         self._client.track_span(
             kind=self.kind,
             name=self.name,
@@ -193,8 +199,6 @@ def __exit__(self, exc_type, exc_val, exc_tb):
             conversation_id=self.conversation_id,
             context=self.context,
         )
-        if self._span_token is not None:
-            _reset_span_context(self._span_token)
         return False
 
     async def __aenter__(self):

wildedge/integrations/openai.py

@@ -41,13 +41,21 @@ def source_from_base_url(base_url: str | None) -> str:
     return SOURCE_BY_HOSTNAME.get(hostname or "", hostname or "openai")
 
 
+def _msg_role(m) -> str | None:
+    return m.get("role") if isinstance(m, dict) else getattr(m, "role", None)
+
+
+def _msg_content(m) -> str | None:
+    return m.get("content") if isinstance(m, dict) else getattr(m, "content", None)
+
+
 def build_input_meta(messages: list, tokens_in: int | None) -> TextInputMeta | None:
     if not messages:
         return None
-    last_user = next((m for m in reversed(messages) if m.get("role") == "user"), None)
+    last_user = next((m for m in reversed(messages) if _msg_role(m) == "user"), None)
     if not last_user:
         return None
-    content = last_user.get("content", "")
+    content = _msg_content(last_user) or ""
     if not isinstance(content, str) or not content:
         return None
     return TextInputMeta(