fixing tracing

agent/core/session.py

@@ -1,13 +1,9 @@
 import asyncio
-import os
 import uuid
 from dataclasses import dataclass
 from enum import Enum
 from typing import Any, Optional
 
-import litellm
-from lmnr import Laminar, LaminarLiteLLMCallback
-
 from agent.config import Config
 from agent.context_manager.manager import ContextManager
 
@@ -38,16 +34,6 @@ class Session:
         event_queue: asyncio.Queue,
         config: Config | None = None,
     ):
-        # Initialize Laminar if API key is present
-        lmnr_api_key = os.environ.get("LMNR_API_KEY")
-        if lmnr_api_key:
-            try:
-                Laminar.initialize(project_api_key=lmnr_api_key)
-                litellm.callbacks = [LaminarLiteLLMCallback()]
-                print("✅ Laminar initialized")
-            except Exception as e:
-                print(f"⚠️ Failed to initialize Laminar: {e}")
-
         self.context_manager = ContextManager()
         self.event_queue = event_queue
         self.session_id = str(uuid.uuid4())

agent/main.py

@@ -3,10 +3,14 @@ Interactive CLI chat with the agent
 """
 
 import asyncio
+import os
 from dataclasses import dataclass
 from pathlib import Path
 from typing import Any, Optional
 
+import litellm
+from lmnr import Laminar, LaminarLiteLLMCallback
+
 from agent.config import load_config
 from agent.core.agent_loop import submission_loop
 from agent.core.session import OpType
@@ -94,6 +98,15 @@ async def main():
     print("=" * 60)
     print("Type your messages below. Type 'exit', 'quit', or '/quit' to end.\n")
 
+    lmnr_api_key = os.environ.get("LMNR_API_KEY")
+    if lmnr_api_key:
+        try:
+            Laminar.initialize(project_api_key=lmnr_api_key)
+            litellm.callbacks = [LaminarLiteLLMCallback()]
+            print("✅ Laminar initialized")
+        except Exception as e:
+            print(f"⚠️ Failed to initialize Laminar: {e}")
+
     # Create queues for communication
     submission_queue = asyncio.Queue()
     event_queue = asyncio.Queue()

eval/hf_agent_connector.py

@@ -5,6 +5,8 @@ import sys
 from pathlib import Path
 from typing import Any
 
+from lmnr import observe
+
 from agent.config import Config, load_config
 from agent.core.agent_loop import Handlers
 from agent.core.session import Session
@@ -38,6 +40,7 @@ class AgentResponseGenerator:
         """Expose the agent model name for downstream logging."""
         return self.config.model_name
 
+    @observe(name="eval_run")
     async def run(self, prompt: str) -> str:
         """
         Execute the agent loop for a single prompt and return the assistant reply.

eval/solvers.py

@@ -10,9 +10,11 @@ import os
 import tempfile
 from typing import Callable, Dict, List, Sequence
 
+import litellm
 from inspect_ai.model import ChatMessageAssistant, ModelOutput
 from inspect_ai.solver import Solver, solver
 from inspect_ai.solver._task_state import TaskState
+from lmnr import Laminar, LaminarLiteLLMCallback
 
 from eval.hf_agent_connector import AgentResponseGenerator
 
@@ -37,6 +39,11 @@ def hf_agent(
     config_path: str = "agent/config_mcp_example.json",
     max_iterations: int = 10,
 ) -> Solver:
+    # init lmnr for observability
+    Laminar.initialize(project_api_key=os.environ.get("LMNR_API_KEY"))
+    litellm.callbacks = [LaminarLiteLLMCallback()]
+    print("✅ Laminar initialized")
+
     runner = AgentResponseGenerator(
         config_path=config_path,
         max_iterations=max_iterations,