adding observability

agent/core/agent_loop.py

@@ -6,6 +6,7 @@ import asyncio
 import json
 
 from litellm import ChatCompletionMessageToolCall, Message, ModelResponse, acompletion
+from lmnr import observe
 
 from agent.config import Config
 from agent.core.session import Event, OpType, Session
@@ -18,8 +19,20 @@ class Handlers:
     """Handler functions for each operation type"""
 
     @staticmethod
-    async def run_agent(session: Session, text: str, max_iterations: int = 10) -> None:
-        """Handle user input (like user_input_or_turn in codex.rs:1291)"""
+    @observe(name="run_agent")
+    async def run_agent(
+        session: Session, text: str, max_iterations: int = 10
+    ) -> str | None:
+        """
+        Handle user input (like user_input_or_turn in codex.rs:1291)
+        Returns the final assistant response content, if any.
+        """
+        # Set session ID for this trace
+        if hasattr(session, "session_id"):
+            from lmnr import Laminar
+
+            Laminar.set_trace_session_id(session_id=session.session_id)
+
         # Add user message to history
         user_msg = Message(role="user", content=text)
         session.context_manager.add_message(user_msg)
@@ -31,6 +44,8 @@ class Handlers:
 
         # Agentic loop - continue until model doesn't call tools or max iterations is reached
         iteration = 0
+        final_response = None
+
         while iteration < max_iterations:
             messages = session.context_manager.get_messages()
             tools = session.tool_router.get_tool_specs_for_llm()
@@ -60,6 +75,7 @@ class Handlers:
                                 data={"content": content},
                             )
                         )
+                        final_response = content
                     break
 
                 # Add assistant message with tool calls to history
@@ -118,7 +134,7 @@ class Handlers:
                 await session.send_event(
                     Event(
                         event_type="error",
-                        data={"error": str(e + "\n" + traceback.format_exc())},
+                        data={"error": str(e) + "\n" + traceback.format_exc()},
                     )
                 )
                 break
@@ -129,6 +145,7 @@ class Handlers:
                 data={"history_size": len(session.context_manager.items)},
             )
         )
+        return final_response
 
     @staticmethod
     async def interrupt(session: Session) -> None:
@@ -202,6 +219,7 @@ async def process_submission(session: Session, submission) -> bool:
     return True
 
 
+@observe(name="submission_loop")
 async def submission_loop(
     submission_queue: asyncio.Queue,
     event_queue: asyncio.Queue,

agent/core/session.py

@@ -1,8 +1,13 @@
 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
 
@@ -33,8 +38,19 @@ 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())
         self.config = config or Config(
             model_name="anthropic/claude-sonnet-4-5-20250929",
             tools=[],

agent/core/tools.py

@@ -9,6 +9,7 @@ from dataclasses import dataclass
 from typing import Any, Awaitable, Callable, Optional
 
 from fastmcp import Client
+from lmnr import observe
 from mcp.types import EmbeddedResource, ImageContent, TextContent
 
 from agent.config import MCPServerConfig
@@ -142,6 +143,7 @@ class ToolRouter:
             await self.mcp_client.__aexit__(exc_type, exc, tb)
             self._mcp_initialized = False
 
+    @observe(name="call_tool")
     async def call_tool(
         self, tool_name: str, arguments: dict[str, Any]
     ) -> tuple[str, bool]:

agent/main.py

@@ -160,6 +160,7 @@ async def main():
                     op_type=OpType.USER_INPUT, data={"text": user_input}
                 ),
             )
+            print(f"Main submitting: {submission.operation.op_type}")
             await submission_queue.put(submission)
 
     except KeyboardInterrupt:

pyproject.toml

@@ -16,4 +16,7 @@ dependencies = [
     "huggingface-hub>=1.0.1",
     "fastmcp>=2.4.0",
     "inspect-ai>=0.3.149",
+    "lmnr[all]>=0.7.23",
+    "transformers>=2.3.0",
+    "torch>=2.9.1",
 ]