Isolate usage tracking and unify trace serialization

Author: haasonsaas

micro_agent/agent.py

@@ -12,14 +12,19 @@ class MicroAgent(dspy.Module):
     The "agent framework": ~100 LOC.
     Plan -> (optional) tool -> observe -> loop -> finalize.
     """
-    def __init__(self, max_steps: int = 6, use_tool_calls: bool | None = None):
+    def __init__(self, max_steps: int = 6, use_tool_calls: bool | None = None, use_global_trace: bool | None = None):
         super().__init__()
         # Use LM directly for robust JSON handling across providers.
         self.lm = dspy.settings.lm
         self.finalize = None  # fallback finalize handled via LM prompt
         self._tool_list = [t.spec() for t in TOOLS.values()]
         self.max_steps = max_steps
         self._provider = self._infer_provider(self.lm)
+        if isinstance(use_global_trace, bool):
+            self._use_global_trace = use_global_trace
+        else:
+            env_gt = os.getenv("MICRO_AGENT_USE_GLOBAL_TRACE")
+            self._use_global_trace = env_gt.strip().lower() not in {"0", "false", "no", "off"} if env_gt else True
         # Determine function-calls mode
         env_override = os.getenv("USE_TOOL_CALLS")
         if isinstance(use_tool_calls, bool):
@@ -205,17 +210,18 @@ def _accumulate_usage(input_text: str = "", output_text: str = ""):
             in_tok = 0
             out_tok = 0
             cost = 0.0
-            try:
-                for _, _, out in dspy.settings.trace[-1:]:
-                    usage = getattr(out, "usage", None) or {}
-                    nonlocal total_cost, total_in_tokens, total_out_tokens
-                    c = getattr(out, "cost", None)
-                    if c is not None:
-                        cost += float(c or 0)
-                    in_tok += int(usage.get("input_tokens", 0) or 0)
-                    out_tok += int(usage.get("output_tokens", 0) or 0)
-            except Exception:
-                pass
+            if self._use_global_trace:
+                try:
+                    for _, _, out in dspy.settings.trace[-1:]:
+                        usage = getattr(out, "usage", None) or {}
+                        nonlocal total_cost, total_in_tokens, total_out_tokens
+                        c = getattr(out, "cost", None)
+                        if c is not None:
+                            cost += float(c or 0)
+                        in_tok += int(usage.get("input_tokens", 0) or 0)
+                        out_tok += int(usage.get("output_tokens", 0) or 0)
+                except Exception:
+                    pass
             if in_tok or out_tok or cost:
                 total_cost += cost
                 total_in_tokens += in_tok

micro_agent/runtime.py

@@ -13,6 +13,13 @@
 def _get_traces_dir() -> str:
     return os.getenv("TRACES_DIR", TRACES_DIR)
 
+def to_jsonable(obj: Any) -> Any:
+    """Coerce arbitrary objects to JSON-serializable structures."""
+    try:
+        return json.loads(json.dumps(obj, ensure_ascii=False, default=str))
+    except Exception:
+        return str(obj)
+
 class Step(TypedDict):
     tool: str
     args: Dict[str, Any]
@@ -44,9 +51,10 @@ def dump_trace(trace_id: str, question: str, steps: List[Step], answer: str, *,
         rec["cost_usd"] = float(cost_usd)
     traces_dir = _get_traces_dir()
     os.makedirs(traces_dir, exist_ok=True)
+    rec_jsonable = to_jsonable(rec)
     path = os.path.join(traces_dir, f"{trace_id}.jsonl")
     with open(path, "a", encoding="utf-8") as f:
-        f.write(json.dumps(rec, ensure_ascii=False, default=str) + "\n")
+        f.write(json.dumps(rec_jsonable, ensure_ascii=False) + "\n")
     return path
 
 def extract_json_block(text: str) -> str:

micro_agent/server.py

@@ -1,22 +1,28 @@
 from __future__ import annotations
 from fastapi import FastAPI, HTTPException
 from fastapi.middleware.cors import CORSMiddleware
-from fastapi.encoders import jsonable_encoder
 import os, json, re
 from threading import Lock
 from pydantic import BaseModel
 from .costs import estimate_prediction_cost
 from importlib.metadata import version as _pkg_version, PackageNotFoundError
 from .config import configure_lm
 from .agent import MicroAgent
-from .runtime import dump_trace, new_trace_id
+from .runtime import dump_trace, new_trace_id, to_jsonable
 from .logging_setup import setup_logging
 
 app = FastAPI(title="DSPy Micro Agent")
+origins_env = os.getenv("MICRO_AGENT_CORS_ORIGINS", "*").strip()
+if origins_env == "*":
+    allow_origins = ["*"]
+    allow_credentials = False
+else:
+    allow_origins = [o.strip() for o in origins_env.split(",") if o.strip()]
+    allow_credentials = os.getenv("MICRO_AGENT_CORS_CREDENTIALS", "0").strip().lower() in {"1", "true", "yes", "on"}
 app.add_middleware(
     CORSMiddleware,
-    allow_origins=["*"],
-    allow_credentials=True,
+    allow_origins=allow_origins,
+    allow_credentials=allow_credentials,
     allow_methods=["*"],
     allow_headers=["*"],
 )
@@ -63,7 +69,9 @@ def ask(req: AskRequest):
         raise HTTPException(status_code=400, detail="max_steps must be between 1 and 20")
 
     def _call_agent():
-        agent = _agent if req.use_tool_calls is None and req.max_steps == _agent.max_steps else MicroAgent(max_steps=req.max_steps, use_tool_calls=req.use_tool_calls)
+        if _serialize and req.use_tool_calls is None and req.max_steps == _agent.max_steps:
+            return _agent(question)
+        agent = MicroAgent(max_steps=req.max_steps, use_tool_calls=req.use_tool_calls, use_global_trace=_serialize)
         return agent(question)
 
     if _serialize:
@@ -76,7 +84,7 @@ def _call_agent():
     usage = getattr(pred, "usage", {}) or {}
     est = estimate_prediction_cost(question, pred.trace, pred.answer, usage)
     path = dump_trace(trace_id, question, pred.trace, pred.answer, usage=usage, cost_usd=est.get("cost_usd"))
-    steps = jsonable_encoder(pred.trace)
+    steps = to_jsonable(pred.trace)
     return AskResponse(answer=pred.answer, trace_id=trace_id, trace_path=path, steps=steps, usage=usage, cost_usd=est.get("cost_usd"))
 
 @app.get("/trace/{trace_id}")

tests/test_regressions.py

@@ -46,6 +46,7 @@ def test_dump_trace_serializes_non_json(tmp_path, monkeypatch):
     with open(path, "r", encoding="utf-8") as f:
         rec = json.loads(f.readline())
     assert rec["steps"][0]["observation"]["when"].startswith("2020-01-01")
+    assert rec["steps"] == runtime.to_jsonable(steps)
 
 
 def test_result_magnitude_limit():