tracellm-cli 0.1.0__py3-none-any.whl

tracellm/integrations/langchain.py

@@ -0,0 +1,186 @@
+import time
+import uuid
+from datetime import datetime, timezone
+from typing import Any, Optional
+
+from langchain_core.callbacks import BaseCallbackHandler
+from langchain_core.messages import BaseMessage
+
+from tracellm.utils import build_tool_step, estimate_tokens, console
+from tracellm.db import save_trace_payload
+
+
+def _get_trace_context() -> dict[str, Any] | None:
+    try:
+        from tracellm.tracer import _current_trace_context
+        return _current_trace_context.get()
+    except ImportError:
+        return None
+
+
+def _finalize_lc_trace(trace_data: dict[str, Any]) -> None:
+    ctx = _get_trace_context() or {}
+    trace_data.setdefault("project_id", ctx.get("project_id") or "default")
+    trace_data.setdefault("project_name", ctx.get("project_name"))
+    trace_data.setdefault("environment", ctx.get("environment") or "development")
+    trace_data.setdefault("api_key", ctx.get("api_key"))
+    try:
+        save_trace_payload(trace_data)
+    except Exception as e:
+        console.print(f"[yellow]LangChain trace persist skipped: {e}[/yellow]")
+
+
+class TracellmCallbackHandler(BaseCallbackHandler):
+    def __init__(self) -> None:
+        super().__init__()
+        self.trace_id = str(uuid.uuid4())
+        self.steps: list[dict[str, Any]] = []
+        self.start_time = time.perf_counter()
+        self.start_dt = datetime.now(timezone.utc)
+        self._chain_stack: list[dict[str, Any]] = []
+        self._retry_count = 0
+        self._error: str | None = None
+        self._llm_inputs: list[str] = []
+        self._llm_outputs: list[str] = []
+
+    @property
+    def always_verbose(self) -> bool:
+        return True
+
+    def on_llm_start(self, serialized: dict[str, Any], prompts: list[str], **kwargs: Any) -> None:
+        self._llm_inputs.append(str(prompts))
+
+    def on_llm_end(self, response: Any, **kwargs: Any) -> None:
+        content = ""
+        try:
+            if hasattr(response, "generations"):
+                content = str(response.generations)
+            elif hasattr(response, "text"):
+                content = response.text
+        except Exception:
+            content = str(response)
+        self._llm_outputs.append(content[:500])
+
+    def on_llm_error(self, error: Exception | KeyboardInterrupt, **kwargs: Any) -> None:
+        self._error = str(error)
+
+    def on_chain_start(self, serialized: dict[str, Any], inputs: dict[str, Any], **kwargs: Any) -> None:
+        step = {
+            "step_id": str(uuid.uuid4()),
+            "tool_name": serialized.get("id", ["unknown"])[-1],
+            "input": {"inputs": inputs},
+            "output": {},
+            "duration": 0.0,
+            "success": True,
+            "timestamp": datetime.now(timezone.utc).isoformat(),
+            "_start": time.perf_counter(),
+        }
+        self._chain_stack.append(step)
+
+    def on_chain_end(self, outputs: dict[str, Any], **kwargs: Any) -> None:
+        if self._chain_stack:
+            step = self._chain_stack.pop()
+            step["duration"] = round((time.perf_counter() - step["_start"]) * 1000, 2)
+            step["output"] = {"outputs": outputs}
+            del step["_start"]
+            self.steps.append(step)
+
+    def on_chain_error(self, error: Exception | KeyboardInterrupt, **kwargs: Any) -> None:
+        if self._chain_stack:
+            step = self._chain_stack.pop()
+            step["duration"] = round((time.perf_counter() - step["_start"]) * 1000, 2)
+            step["output"] = {"error": str(error)}
+            step["success"] = False
+            del step["_start"]
+            self.steps.append(step)
+        self._error = str(error)
+
+    def on_tool_start(self, serialized: dict[str, Any], input_str: str, **kwargs: Any) -> None:
+        step = {
+            "step_id": str(uuid.uuid4()),
+            "tool_name": serialized.get("name", "unknown_tool"),
+            "input": {"input": input_str},
+            "output": {},
+            "duration": 0.0,
+            "success": True,
+            "timestamp": datetime.now(timezone.utc).isoformat(),
+            "_start": time.perf_counter(),
+        }
+        self._chain_stack.append(step)
+
+    def on_tool_end(self, output: str, **kwargs: Any) -> None:
+        if self._chain_stack:
+            step = self._chain_stack.pop()
+            step["duration"] = round((time.perf_counter() - step["_start"]) * 1000, 2)
+            step["output"] = {"output": output[:500]}
+            del step["_start"]
+            self.steps.append(step)
+
+    def on_tool_error(self, error: Exception | KeyboardInterrupt, **kwargs: Any) -> None:
+        if self._chain_stack:
+            step = self._chain_stack.pop()
+            step["duration"] = round((time.perf_counter() - step["_start"]) * 1000, 2)
+            step["output"] = {"error": str(error)}
+            step["success"] = False
+            del step["_start"]
+            self.steps.append(step)
+        self._retry_count += 1
+
+    def on_retriever_start(self, serialized: dict[str, Any], query: str, **kwargs: Any) -> None:
+        step = {
+            "step_id": str(uuid.uuid4()),
+            "tool_name": "retriever",
+            "input": {"query": query},
+            "output": {},
+            "duration": 0.0,
+            "success": True,
+            "timestamp": datetime.now(timezone.utc).isoformat(),
+            "_start": time.perf_counter(),
+        }
+        self._chain_stack.append(step)
+
+    def on_retriever_end(self, documents: list[Any], **kwargs: Any) -> None:
+        if self._chain_stack:
+            step = self._chain_stack.pop()
+            step["duration"] = round((time.perf_counter() - step["_start"]) * 1000, 2)
+            doc_previews = [str(doc)[:200] for doc in documents[:5]]
+            step["output"] = {"documents": doc_previews, "count": len(documents)}
+            del step["_start"]
+            self.steps.append(step)
+
+    def on_retriever_error(self, error: Exception | KeyboardInterrupt, **kwargs: Any) -> None:
+        if self._chain_stack:
+            step = self._chain_stack.pop()
+            step["duration"] = round((time.perf_counter() - step["_start"]) * 1000, 2)
+            step["output"] = {"error": str(error)}
+            step["success"] = False
+            del step["_start"]
+            self.steps.append(step)
+
+    def flush_trace(
+        self,
+        prompt: str | None = None,
+        response: str | None = None,
+        status: str = "success",
+    ) -> dict[str, Any]:
+        latency = (time.perf_counter() - self.start_time) * 1000
+        total_input = " ".join(self._llm_inputs) if self._llm_inputs else (prompt or "")
+        total_output = " ".join(self._llm_outputs) if self._llm_outputs else (response or "")
+
+        trace_data = {
+            "trace_id": self.trace_id,
+            "prompt": total_input,
+            "response": total_output,
+            "latency": round(latency, 2),
+            "token_count": estimate_tokens(total_input, total_output),
+            "model_name": "langchain",
+            "status": status if not self._error else "failed",
+            "steps": self.steps,
+            "retry_count": self._retry_count,
+            "failure_reason": self._error,
+            "slow_request": latency >= 1500.0,
+            "created_at": self.start_dt.isoformat(),
+            "updated_at": datetime.now(timezone.utc).isoformat(),
+        }
+        _finalize_lc_trace(trace_data)
+        return trace_data

tracellm/integrations/openai.py

@@ -0,0 +1,234 @@
+import functools
+import time
+import uuid
+from datetime import datetime, timezone
+from typing import Any, Iterator
+
+from openai import OpenAI, Stream
+from openai.types.chat import ChatCompletion, ChatCompletionChunk
+
+from tracellm.utils import (
+    console,
+    estimate_tokens,
+    build_tool_step,
+)
+from tracellm.db import save_trace_payload
+
+
+def _get_openai_trace_context() -> dict[str, Any] | None:
+    try:
+        from tracellm.tracer import _current_trace_context
+        return _current_trace_context.get()
+    except ImportError:
+        return None
+
+
+def _build_openai_step(
+    tool_name: str,
+    input_data: dict[str, Any],
+    output_data: dict[str, Any],
+    duration: float,
+    success: bool = True,
+) -> dict[str, Any]:
+    return build_tool_step(tool_name, input_data, output_data, duration, success)
+
+
+def _finalize_trace(
+    trace_id: str,
+    prompt: str,
+    response: str,
+    model: str,
+    latency: float,
+    token_count: int,
+    steps: list[dict[str, Any]],
+    status: str,
+    error: str | None = None,
+    retry_count: int = 0,
+) -> dict[str, Any]:
+    now = datetime.now(timezone.utc)
+    ctx = _get_openai_trace_context()
+    trace_data = {
+        "trace_id": trace_id,
+        "prompt": prompt,
+        "response": response,
+        "latency": round(latency, 2),
+        "token_count": token_count,
+        "model_name": model,
+        "project_id": (ctx.get("project_id") if ctx else None) or "default",
+        "project_name": (ctx.get("project_name") if ctx else None) or None,
+        "api_key": (ctx.get("api_key") if ctx else None) or None,
+        "environment": (ctx.get("environment") if ctx else None) or "development",
+        "status": status,
+        "steps": steps,
+        "retry_count": retry_count,
+        "failure_reason": error,
+        "slow_request": latency >= 1500.0,
+        "created_at": now.isoformat(),
+        "updated_at": now.isoformat(),
+    }
+    try:
+        save_trace_payload(trace_data)
+    except Exception as e:
+        console.print(f"[yellow]Trace persist skipped: {e}[/yellow]")
+    return trace_data
+
+
+def _extract_token_usage(completion: ChatCompletion) -> dict[str, int]:
+    usage = completion.usage
+    if usage:
+        return {
+            "prompt_tokens": usage.prompt_tokens,
+            "completion_tokens": usage.completion_tokens,
+            "total_tokens": usage.total_tokens,
+        }
+    return {}
+
+
+def _estimate_streaming_tokens(chunks: list[ChatCompletionChunk]) -> dict[str, int]:
+    total = 0
+    for chunk in chunks:
+        if chunk.choices and chunk.choices[0].delta.content:
+            total += 1
+    return {"total_tokens": total}
+
+
+def wrap_openai(client: OpenAI) -> OpenAI:
+    original_create = client.chat.completions.create
+
+    @functools.wraps(original_create)
+    def traced_create(*args: Any, **kwargs: Any) -> Any:
+        return _traced_chat_completion(client, kwargs, original_create)
+
+    client.chat.completions.create = traced_create
+    return client
+
+
+def _traced_chat_completion(client: OpenAI, kwargs: dict[str, Any], original_create: Any) -> Any:
+    trace_id = str(uuid.uuid4())
+    started_at = time.perf_counter()
+    start_dt = datetime.now(timezone.utc)
+
+    messages = kwargs.get("messages", [])
+    model = kwargs.get("model", "unknown")
+    stream = kwargs.get("stream", False)
+
+    prompt_text = str(messages)
+    steps: list[dict[str, Any]] = []
+    error: str | None = None
+    retry_count = 0
+    max_retries = kwargs.get("max_retries", 0) or 0
+
+    max_attempts = max_retries + 1
+    last_exception: Exception | None = None
+
+    for attempt in range(max_attempts):
+        step_start = time.perf_counter()
+        try:
+            if attempt > 0:
+                retry_count += 1
+
+            result = original_create(*client.chat.completions._get_original_args() if hasattr(client.chat.completions, '_get_original_args') else [], **kwargs)
+
+            if stream:
+                collected_chunks: list[ChatCompletionChunk] = []
+                full_content = ""
+
+                for chunk in result:
+                    collected_chunks.append(chunk)
+                    if chunk.choices and chunk.choices[0].delta.content:
+                        full_content += chunk.choices[0].delta.content
+                    yield chunk
+
+                step_duration = (time.perf_counter() - step_start) * 1000
+                token_data = _estimate_streaming_tokens(collected_chunks)
+                token_count = sum(token_data.values())
+                steps.append(_build_openai_step(
+                    "openai_chat_stream",
+                    {"model": model, "messages": messages},
+                    {"content_preview": full_content[:200], "chunks": len(collected_chunks)},
+                    step_duration,
+                ))
+
+                latency = (time.perf_counter() - started_at) * 1000
+                _finalize_trace(
+                    trace_id=trace_id,
+                    prompt=prompt_text,
+                    response=full_content,
+                    model=model,
+                    latency=latency,
+                    token_count=token_count,
+                    steps=steps,
+                    status="success",
+                    retry_count=retry_count,
+                )
+                return
+
+            step_duration = (time.perf_counter() - step_start) * 1000
+
+            token_usage = _extract_token_usage(result)
+            token_count = token_usage.get("total_tokens", estimate_tokens(str(messages), str(result.choices[0].message.content if result.choices else "")))
+
+            steps.append(_build_openai_step(
+                "openai_chat",
+                {"model": model, "messages": messages},
+                {
+                    "content": result.choices[0].message.content if result.choices else "",
+                    "finish_reason": result.choices[0].finish_reason if result.choices else None,
+                    "usage": token_usage,
+                },
+                step_duration,
+            ))
+
+            latency = (time.perf_counter() - started_at) * 1000
+            response_text = result.choices[0].message.content if result.choices else ""
+            _finalize_trace(
+                trace_id=trace_id,
+                prompt=prompt_text,
+                response=response_text,
+                model=model,
+                latency=latency,
+                token_count=token_count,
+                steps=steps,
+                status="success",
+                retry_count=retry_count,
+            )
+            return result
+
+        except Exception as e:
+            last_exception = e
+            step_duration = (time.perf_counter() - step_start) * 1000
+            error = str(e)
+            steps.append(_build_openai_step(
+                f"openai_chat_retry_{attempt + 1}" if attempt < max_attempts - 1 else "openai_chat",
+                {"model": model, "messages": messages, "attempt": attempt + 1},
+                {"error": error},
+                step_duration,
+                success=(attempt == max_attempts - 1),
+            ))
+            if attempt < max_attempts - 1:
+                time.sleep(min(0.5 * (2 ** attempt), 5.0))
+            else:
+                break
+
+    latency = (time.perf_counter() - started_at) * 1000
+    _finalize_trace(
+        trace_id=trace_id,
+        prompt=prompt_text,
+        response="",
+        model=model,
+        latency=latency,
+        token_count=0,
+        steps=steps,
+        status="failed",
+        error=error,
+        retry_count=retry_count,
+    )
+    if last_exception:
+        raise last_exception
+    return None
+
+
+class TraceOpenAI(OpenAI):
+    def __init__(self, *args: Any, **kwargs: Any) -> None:
+        super().__init__(*args, **kwargs)
+        wrap_openai(self)

tracellm/integrations/tool_tracer.py

@@ -0,0 +1,151 @@
+import asyncio
+import functools
+import inspect
+import time
+import uuid
+from datetime import datetime, timezone
+from typing import Any, Callable, TypeVar
+
+from tracellm.utils import console
+from tracellm.db import save_trace_payload
+
+F = TypeVar("F", bound=Callable[..., Any])
+
+
+def trace_tool(
+    name: str | None = None,
+    max_retries: int = 0,
+    capture_input: bool = True,
+    capture_output: bool = True,
+) -> Callable[[F], F]:
+    def decorator(func: F) -> F:
+        tool_name = name or func.__name__
+        retry_count = 0
+
+        if inspect.iscoroutinefunction(func):
+
+            @functools.wraps(func)
+            async def async_wrapper(*args: Any, **kwargs: Any) -> Any:
+                nonlocal retry_count
+                step_id = str(uuid.uuid4())
+                started_at = time.perf_counter()
+                last_exception: Exception | None = None
+
+                for attempt in range(max_retries + 1):
+                    try:
+                        if attempt > 0:
+                            retry_count += 1
+                            console.print(f"  [yellow]retry {attempt}/{max_retries}[/yellow] {tool_name}")
+
+                        result = await func(*args, **kwargs)
+                        duration = (time.perf_counter() - started_at) * 1000
+
+                        step = {
+                            "step_id": step_id,
+                            "tool_name": tool_name,
+                            "input": _capture_args(func, args, kwargs) if capture_input else {},
+                            "output": _capture_output(result) if capture_output else {},
+                            "duration": round(duration, 2),
+                            "success": True,
+                            "timestamp": datetime.now(timezone.utc).isoformat(),
+                        }
+                        _try_append_step(step)
+                        return result
+
+                    except Exception as e:
+                        last_exception = e
+                        duration = (time.perf_counter() - started_at) * 1000
+                        if attempt < max_retries:
+                            wait = min(0.5 * (2 ** attempt), 5.0)
+                            await asyncio.sleep(wait)
+                        else:
+                            step = {
+                                "step_id": step_id,
+                                "tool_name": tool_name,
+                                "input": _capture_args(func, args, kwargs) if capture_input else {},
+                                "output": {"error": str(e)},
+                                "duration": round(duration, 2),
+                                "success": False,
+                                "timestamp": datetime.now(timezone.utc).isoformat(),
+                            }
+                            _try_append_step(step)
+                            raise
+
+            return async_wrapper  # type: ignore
+
+        else:
+
+            @functools.wraps(func)
+            def sync_wrapper(*args: Any, **kwargs: Any) -> Any:
+                nonlocal retry_count
+                step_id = str(uuid.uuid4())
+                started_at = time.perf_counter()
+                last_exception: Exception | None = None
+
+                for attempt in range(max_retries + 1):
+                    try:
+                        if attempt > 0:
+                            retry_count += 1
+                            console.print(f"  [yellow]retry {attempt}/{max_retries}[/yellow] {tool_name}")
+
+                        result = func(*args, **kwargs)
+                        duration = (time.perf_counter() - started_at) * 1000
+
+                        step = {
+                            "step_id": step_id,
+                            "tool_name": tool_name,
+                            "input": _capture_args(func, args, kwargs) if capture_input else {},
+                            "output": _capture_output(result) if capture_output else {},
+                            "duration": round(duration, 2),
+                            "success": True,
+                            "timestamp": datetime.now(timezone.utc).isoformat(),
+                        }
+                        _try_append_step(step)
+                        return result
+
+                    except Exception as e:
+                        last_exception = e
+                        duration = (time.perf_counter() - started_at) * 1000
+                        if attempt < max_retries:
+                            time.sleep(min(0.5 * (2 ** attempt), 5.0))
+                        else:
+                            step = {
+                                "step_id": step_id,
+                                "tool_name": tool_name,
+                                "input": _capture_args(func, args, kwargs) if capture_input else {},
+                                "output": {"error": str(e)},
+                                "duration": round(duration, 2),
+                                "success": False,
+                                "timestamp": datetime.now(timezone.utc).isoformat(),
+                            }
+                            _try_append_step(step)
+                            raise
+
+            return sync_wrapper  # type: ignore
+
+    return decorator
+
+
+def _capture_args(func: Callable, args: tuple[Any, ...], kwargs: dict[str, Any]) -> dict[str, Any]:
+    try:
+        sig = inspect.signature(func)
+        bound = sig.bind(*args, **kwargs)
+        bound.apply_defaults()
+        return {k: str(v)[:500] for k, v in bound.arguments.items()}
+    except Exception:
+        return {"args": str(args)[:500], "kwargs": str(kwargs)[:500]}
+
+
+def _capture_output(result: Any) -> dict[str, Any]:
+    return {"result": str(result)[:500]}
+
+
+def _try_append_step(step: dict[str, Any]) -> None:
+    try:
+        from tracellm.tracer import _current_trace_context
+        ctx = _current_trace_context.get()
+        if ctx is not None:
+            steps = ctx.setdefault("collected_steps", [])
+            steps.append(step)
+    except Exception:
+        pass

tracellm/mascot.py

@@ -0,0 +1,49 @@
+"""TraceLLM dinosaur mascot."""
+
+from __future__ import annotations
+
+from enum import Enum
+
+from rich.panel import Panel
+from rich.text import Text
+
+
+class MascotState(Enum):
+    IDLE = "idle"
+    LOADING = "loading"
+    SUCCESS = "success"
+    WARNING = "warning"
+
+
+DINOSAUR = """\
+      __
+     / _)
+ .-^^^-/ /
+__/       /
+<__.|_|-|_|"""
+
+
+_STYLE: dict[MascotState, str] = {
+    MascotState.IDLE: "bright_black",
+    MascotState.LOADING: "cyan",
+    MascotState.SUCCESS: "green",
+    MascotState.WARNING: "yellow",
+}
+
+
+def render(state: MascotState = MascotState.IDLE) -> Text:
+    """Full ASCII dinosaur art with state-based styling."""
+    return Text(DINOSAUR, style=_STYLE[state])
+
+
+def header(title: str, state: MascotState = MascotState.IDLE) -> Panel:
+    """Compact header Panel with mascot prefix."""
+    dino = Text("🦖 ", style=_STYLE[state])
+    title_text = Text.assemble(dino, Text(title, style="bold white"))
+    return Panel("", title=title_text, border_style="bright_black")
+
+
+def message(text: str, state: MascotState = MascotState.IDLE) -> Text:
+    """One-line mascot message (e.g. \"🦖 Trace complete\")."""
+    dino = Text("🦖 ", style=_STYLE[state])
+    return Text.assemble(dino, Text(text, style=_STYLE[state]))