noesium 0.1.0__py3-none-any.whl

noesium/core/tracing/langgraph_hooks.py

@@ -0,0 +1,308 @@
+import time
+from datetime import datetime
+from typing import Any, Dict, List, Optional
+
+from langchain_core.callbacks import BaseCallbackHandler
+from langchain_core.outputs import LLMResult
+
+from noesium.core.tracing.token_tracker import get_token_tracker
+from noesium.core.utils.logging import color_text, get_logger
+
+logger = get_logger(__name__)
+
+
+class NodeLoggingCallback(BaseCallbackHandler):
+    def __init__(self, node_id: Optional[str] = None):
+        self.node_id = node_id
+
+    def _prefix(self) -> str:
+        return f"[{self.node_id}] " if self.node_id else ""
+
+    def on_tool_end(self, output, run_id, parent_run_id, **kwargs):
+        logger.info(color_text(f"{self._prefix()}[TOOL END] output={output}", "cyan", ["dim"]))
+
+    def on_chain_end(self, output, run_id, parent_run_id, **kwargs):
+        logger.info(color_text(f"{self._prefix()}[CHAIN END] output={output}", "blue", ["dim"]))
+
+    def on_llm_end(self, response, run_id, parent_run_id, **kwargs):
+        logger.info(color_text(f"{self._prefix()}[LLM END] response={response}", "magenta", ["dim"]))
+
+    def on_custom_event(self, event_name, payload, **kwargs):
+        logger.info(color_text(f"{self._prefix()}[EVENT] {event_name}: {payload}", "yellow", ["dim"]))
+
+
+class TokenUsageCallback(BaseCallbackHandler):
+    """Enhanced token usage callback that works with LangGraph and custom LLM clients."""
+
+    def __init__(self, model_name: Optional[str] = None, verbose: bool = True):
+        super().__init__()
+        self.total_prompt_tokens = 0
+        self.total_completion_tokens = 0
+        self.model_name = model_name
+        self.verbose = verbose
+        self.session_start = time.time()
+        self.llm_calls = 0
+        self.token_usage_history: List[Dict] = []
+
+        # Track custom events from our LLM clients
+        self._pending_calls: Dict[str, Dict] = {}
+
+    def on_llm_start(
+        self, serialized: Dict, prompts: List[str], run_id: str, parent_run_id: Optional[str] = None, **kwargs
+    ):
+        """Track when LLM calls start"""
+        self.llm_calls += 1
+
+        # Store prompt info for token counting
+        self._pending_calls[run_id] = {"prompts": prompts, "start_time": time.time(), "call_number": self.llm_calls}
+
+        if self.verbose:
+            logger.info(
+                color_text(
+                    f"[TOKEN CALLBACK] LLM call #{self.llm_calls} started (run_id: {run_id[:8]}...)", "magenta", ["dim"]
+                )
+            )
+
+    def on_llm_end(self, response: LLMResult, run_id: str, parent_run_id: Optional[str] = None, **kwargs):
+        """Enhanced token usage tracking with multiple extraction methods"""
+        usage_data = self._extract_token_usage_from_response(response)
+
+        # Fallback to prompt estimation if no usage data available
+        if not usage_data and run_id in self._pending_calls:
+            usage_data = self._estimate_token_usage(self._pending_calls[run_id], response)
+
+        if usage_data:
+            prompt_tokens = usage_data.get("prompt_tokens", 0)
+            completion_tokens = usage_data.get("completion_tokens", 0)
+            total_tokens = usage_data.get("total_tokens", prompt_tokens + completion_tokens)
+
+            # Update totals
+            self.total_prompt_tokens += prompt_tokens
+            self.total_completion_tokens += completion_tokens
+
+            # Store in history
+            call_data = {
+                "run_id": run_id,
+                "timestamp": datetime.now().isoformat(),
+                "prompt_tokens": prompt_tokens,
+                "completion_tokens": completion_tokens,
+                "total_tokens": total_tokens,
+                "model_name": self.model_name or "unknown",
+            }
+            self.token_usage_history.append(call_data)
+
+            if self.verbose:
+                self._log_token_usage(prompt_tokens, completion_tokens, total_tokens, run_id)
+            else:
+                # Even when not verbose, show basic token usage at info level
+                token_log = f"TOKENS: {total_tokens} | langchain | unknown | P:{prompt_tokens} C:{completion_tokens}"
+                logger.info(color_text(token_log, "magenta"))
+
+        # Clean up pending call
+        self._pending_calls.pop(run_id, None)
+
+    def on_custom_event(self, name: str, data: Any, run_id: str, **kwargs):
+        """Handle custom events from our LLM clients for token usage"""
+        if name == "token_usage":
+            self._handle_custom_token_usage(data, run_id)
+
+    def _handle_custom_token_usage(self, data: Dict, run_id: str):
+        """Handle custom token usage events from our LLM clients"""
+        if isinstance(data, dict) and "usage" in data:
+            usage = data["usage"]
+            prompt_tokens = usage.get("prompt_tokens", 0)
+            completion_tokens = usage.get("completion_tokens", 0)
+            total_tokens = usage.get("total_tokens", prompt_tokens + completion_tokens)
+
+            # Update totals
+            self.total_prompt_tokens += prompt_tokens
+            self.total_completion_tokens += completion_tokens
+
+            # Store in history
+            call_data = {
+                "run_id": run_id,
+                "timestamp": datetime.now().isoformat(),
+                "prompt_tokens": prompt_tokens,
+                "completion_tokens": completion_tokens,
+                "total_tokens": total_tokens,
+                "model_name": self.model_name or data.get("model", "unknown"),
+                "source": "custom_event",
+            }
+            self.token_usage_history.append(call_data)
+
+            if self.verbose:
+                self._log_token_usage(prompt_tokens, completion_tokens, total_tokens, run_id)
+            else:
+                # Even when not verbose, show basic token usage at info level
+                model_name = (
+                    data.get("model", "unknown").split("/")[-1]
+                    if "/" in data.get("model", "unknown")
+                    else data.get("model", "unknown")
+                )
+                token_log = (
+                    f"TOKENS: {total_tokens} | custom_event | {model_name} | P:{prompt_tokens} C:{completion_tokens}"
+                )
+                logger.info(color_text(token_log, "magenta"))
+
+    def _extract_token_usage_from_response(self, response: LLMResult) -> Optional[Dict]:
+        """Extract token usage from LangChain LLMResult"""
+        usage = None
+
+        # Method 1: LLMResult llm_output
+        if hasattr(response, "llm_output") and response.llm_output:
+            usage = response.llm_output.get("token_usage") or response.llm_output.get("usage")
+
+        # Method 2: LLMResult response_metadata
+        if not usage and hasattr(response, "response_metadata") and response.response_metadata:
+            usage = response.response_metadata.get("token_usage") or response.response_metadata.get("usage")
+
+        # Method 3: Check generations for usage info
+        if not usage and hasattr(response, "generations") and response.generations:
+            for generation_list in response.generations:
+                for generation in generation_list:
+                    if hasattr(generation, "generation_info") and generation.generation_info:
+                        gen_usage = generation.generation_info.get("token_usage") or generation.generation_info.get(
+                            "usage"
+                        )
+                        if gen_usage:
+                            usage = gen_usage
+                            break
+                if usage:
+                    break
+
+        return usage
+
+    def _estimate_token_usage(self, call_info: Dict, response: LLMResult) -> Dict:
+        """Estimate token usage when actual usage is not available"""
+        try:
+            # Simple estimation: ~4 characters per token for English text
+            prompts = call_info.get("prompts", [])
+            prompt_chars = sum(len(prompt) for prompt in prompts)
+            prompt_tokens = max(1, prompt_chars // 4)
+
+            # Estimate completion tokens from response
+            completion_chars = 0
+            if hasattr(response, "generations"):
+                for generation_list in response.generations:
+                    for generation in generation_list:
+                        if hasattr(generation, "text"):
+                            completion_chars += len(generation.text)
+
+            completion_tokens = max(1, completion_chars // 4)
+
+            return {
+                "prompt_tokens": prompt_tokens,
+                "completion_tokens": completion_tokens,
+                "total_tokens": prompt_tokens + completion_tokens,
+                "estimated": True,
+            }
+        except Exception as e:
+            logger.warning(f"Failed to estimate token usage: {e}")
+            return {"prompt_tokens": 0, "completion_tokens": 0, "total_tokens": 0, "estimated": True, "error": str(e)}
+
+    def _log_token_usage(
+        self, prompt_tokens: int, completion_tokens: int, total_tokens: int, run_id: Optional[str] = None
+    ):
+        """Log token usage with detailed information"""
+        run_info = f" (run_id: {run_id[:8]}...)" if run_id else ""
+        model_info = f" [{self.model_name}]" if self.model_name else ""
+
+        logger.info(color_text(f"[TOKEN USAGE]{model_info}{run_info}", "magenta", ["dim"]))
+        logger.info(color_text(f"  Prompt: {prompt_tokens:,} tokens", None, ["dim"]))
+        logger.info(color_text(f"  Completion: {completion_tokens:,} tokens", None, ["dim"]))
+        logger.info(color_text(f"  Total: {total_tokens:,} tokens", None, ["dim"]))
+
+        # Show session totals
+        session_total = self.total_tokens()
+        logger.info(color_text(f"  Session Total: {session_total:,} tokens", None, ["dim"]))
+
+    def total_tokens(self) -> int:
+        """Get total tokens used in this session"""
+        return self.total_prompt_tokens + self.total_completion_tokens
+
+    def get_session_summary(self) -> Dict:
+        """Get comprehensive session summary including custom LLM client usage"""
+        session_duration = time.time() - self.session_start
+
+        # Get data from global token tracker (custom LLM clients)
+        tracker_stats = get_token_tracker().get_stats()
+
+        # Combine callback and tracker statistics
+        combined_prompt_tokens = self.total_prompt_tokens + tracker_stats.get("total_prompt_tokens", 0)
+        combined_completion_tokens = self.total_completion_tokens + tracker_stats.get("total_completion_tokens", 0)
+        combined_total_tokens = combined_prompt_tokens + combined_completion_tokens
+        combined_calls = self.llm_calls + tracker_stats.get("total_calls", 0)
+
+        return {
+            "session_duration_seconds": session_duration,
+            # Combined stats from both callback and tracker
+            "total_llm_calls": combined_calls,
+            "total_prompt_tokens": combined_prompt_tokens,
+            "total_completion_tokens": combined_completion_tokens,
+            "total_tokens": combined_total_tokens,
+            # Separate stats for debugging
+            "callback_stats": {
+                "llm_calls": self.llm_calls,
+                "prompt_tokens": self.total_prompt_tokens,
+                "completion_tokens": self.total_completion_tokens,
+                "total_tokens": self.total_tokens(),
+            },
+            "tracker_stats": tracker_stats,
+            "model_name": self.model_name,
+            "token_usage_history": self.token_usage_history,
+        }
+
+    def print_session_summary(self):
+        """Print a formatted session summary"""
+        summary = self.get_session_summary()
+
+        logger.info(color_text("\n" + "=" * 50, "blue", ["dim"]))
+        logger.info(color_text("TOKEN USAGE SESSION SUMMARY", "blue", ["dim"]))
+        logger.info(color_text("=" * 50, "blue", ["dim"]))
+        logger.info(color_text(f"Session Duration: {summary['session_duration_seconds']:.2f} seconds", None, ["dim"]))
+        logger.info(color_text(f"Total LLM Calls: {summary['total_llm_calls']}", None, ["dim"]))
+        logger.info(color_text(f"Total Prompt Tokens: {summary['total_prompt_tokens']:,}", None, ["dim"]))
+        logger.info(color_text(f"Total Completion Tokens: {summary['total_completion_tokens']:,}", None, ["dim"]))
+        logger.info(color_text(f"Total Tokens: {summary['total_tokens']:,}", None, ["dim"]))
+
+        # Show breakdown by source
+        callback_stats = summary["callback_stats"]
+        tracker_stats = summary["tracker_stats"]
+
+        if callback_stats["llm_calls"] > 0:
+            logger.info(
+                color_text(
+                    f"  LangChain Calls: {callback_stats['llm_calls']} ({callback_stats['total_tokens']:,} tokens)",
+                    None,
+                    ["dim"],
+                )
+            )
+
+        if tracker_stats.get("total_calls", 0) > 0:
+            logger.info(
+                color_text(
+                    f"  Custom Client Calls: {tracker_stats['total_calls']} ({tracker_stats['total_tokens']:,} tokens)",
+                    None,
+                    ["dim"],
+                )
+            )
+
+        if self.model_name:
+            logger.info(color_text(f"Model: {self.model_name}", None, ["dim"]))
+
+        logger.info(color_text("=" * 50, "blue", ["dim"]))
+
+    def reset_session(self):
+        """Reset all counters for a new session"""
+        self.total_prompt_tokens = 0
+        self.total_completion_tokens = 0
+        self.session_start = time.time()
+        self.llm_calls = 0
+        self.token_usage_history = []
+        self._pending_calls.clear()
+
+        # Also reset the global token tracker
+        get_token_tracker().reset()
+
+        if self.verbose:
+            logger.info(color_text("[TOKEN CALLBACK] Session reset (including custom clients)", "magenta", ["dim"]))

noesium/core/tracing/opik_tracing.py

@@ -0,0 +1,144 @@
+"""
+Opik tracing configuration for observability of LLM communications.
+
+This module configures Opik tracing to monitor and analyze
+all LLM interactions in the Noesium application.
+"""
+
+import logging
+import os
+from typing import Optional
+
+logger = logging.getLogger(__name__)
+
+
+def configure_opik() -> bool:
+    """
+    Configure Opik tracing based on environment variables.
+
+    Environment variables:
+        COGENTS_OPIK_TRACING: Global toggle for Opik tracing (default: false)
+        OPIK_USE_LOCAL: Use local Opik deployment (default: true)
+        OPIK_LOCAL_URL: Local Opik URL (default: http://localhost:5173)
+        OPIK_API_KEY: API key for Comet ML/Opik (only needed for cloud)
+        OPIK_WORKSPACE: Workspace name (optional)
+        OPIK_PROJECT_NAME: Project name for organizing traces
+        OPIK_URL: Custom Opik URL (for cloud deployment)
+        OPIK_TRACING: Enable/disable tracing (default: true if enabled globally)
+
+    Returns:
+        bool: True if Opik was successfully configured, False otherwise
+    """
+    try:
+        # Check global Noesium Opik tracing toggle first
+        cogents_opik_enabled = os.getenv("COGENTS_OPIK_TRACING", "false").lower() == "true"
+
+        if not cogents_opik_enabled:
+            logger.debug("Opik tracing disabled via COGENTS_OPIK_TRACING=false")
+            return False
+
+        # Check if using local deployment
+        use_local = os.getenv("OPIK_USE_LOCAL", "true").lower() == "true"
+
+        # Configuration variables
+        opik_api_key = os.getenv("OPIK_API_KEY")
+        os.getenv("OPIK_WORKSPACE")
+        opik_project = os.getenv("OPIK_PROJECT_NAME", "cogents-llm")
+        opik_tracing = os.getenv("OPIK_TRACING", "true").lower() == "true"
+
+        if not opik_tracing:
+            logger.debug("Opik tracing disabled via OPIK_TRACING=false")
+            return False
+
+        # For cloud deployment, API key is required
+        if not use_local and not opik_api_key:
+            logger.debug("No OPIK_API_KEY found for cloud deployment, Opik tracing disabled")
+            return False
+
+        # Import opik here to avoid import errors if not installed
+        import opik
+
+        # Configure Opik
+        if use_local:
+            # Local deployment configuration
+            local_url = os.getenv("OPIK_LOCAL_URL", "http://localhost:5173")
+            opik.configure(use_local=True)
+            logger.info(f"Opik tracing configured for local deployment at {local_url}, project: {opik_project}")
+        else:
+            # Cloud deployment configuration
+            opik.configure(api_key=opik_api_key, use_local=False)
+            logger.info(f"Opik tracing configured for cloud deployment, project: {opik_project}")
+
+        return True
+
+    except ImportError:
+        logger.warning("Opik package not installed, tracing disabled")
+        return False
+    except Exception as e:
+        logger.error(f"Failed to configure Opik: {e}")
+        return False
+
+
+def is_opik_enabled() -> bool:
+    """
+    Check if Opik tracing is currently enabled.
+
+    Returns:
+        bool: True if Opik tracing is enabled, False otherwise
+    """
+    try:
+        # Check global Noesium Opik tracing toggle first
+        cogents_opik_enabled = os.getenv("COGENTS_OPIK_TRACING", "false").lower() == "true"
+
+        if not cogents_opik_enabled:
+            return False
+
+        import opik
+
+        # Check if Opik is properly configured
+        return opik.get_current_project_name() is not None
+    except (ImportError, AttributeError):
+        return False
+    except Exception:
+        return False
+
+
+def get_opik_project() -> Optional[str]:
+    """
+    Get the current Opik project name.
+
+    Returns:
+        Optional[str]: Project name if configured, None otherwise
+    """
+    try:
+        import opik
+
+        return opik.get_current_project_name()
+    except (ImportError, AttributeError):
+        return None
+    except Exception:
+        return None
+
+
+def create_opik_trace(name: str, input_data: dict = None, metadata: dict = None):
+    """
+    Create a new Opik trace.
+
+    Args:
+        name: Name of the trace
+        input_data: Input data for the trace
+        metadata: Additional metadata
+
+    Returns:
+        Opik trace object or None if tracing is disabled
+    """
+    if not is_opik_enabled():
+        return None
+
+    try:
+        import opik
+
+        return opik.trace(name=name, input=input_data or {}, metadata=metadata or {})
+    except Exception as e:
+        logger.debug(f"Failed to create Opik trace: {e}")
+        return None

noesium/core/tracing/token_tracker.py

@@ -0,0 +1,166 @@
+"""
+Token usage tracking system for custom LLM clients.
+
+This module provides a thread-safe token usage tracker that can be used
+by custom LLM clients to report usage and by callbacks to access statistics.
+"""
+
+import threading
+import time
+from dataclasses import dataclass
+from datetime import datetime
+from typing import Any, Dict, List, Optional
+
+from noesium.core.utils.logging import color_text, get_logger
+
+logger = get_logger(__name__)
+
+
+@dataclass
+class TokenUsage:
+    """Token usage data for a single LLM call."""
+
+    prompt_tokens: int
+    completion_tokens: int
+    total_tokens: int
+    model_name: str
+    timestamp: str
+    run_id: Optional[str] = None
+    call_type: str = "completion"  # completion, structured, vision
+    estimated: bool = False
+
+
+class TokenUsageTracker:
+    """Thread-safe token usage tracker."""
+
+    def __init__(self):
+        self._lock = threading.RLock()  # Use reentrant lock to avoid deadlock
+        self.total_prompt_tokens = 0
+        self.total_completion_tokens = 0
+        self.session_start = time.time()
+        self.call_count = 0
+        self.usage_history: List[TokenUsage] = []
+
+    def record_usage(self, usage: TokenUsage) -> None:
+        """Record token usage from an LLM call."""
+        with self._lock:
+            self.total_prompt_tokens += usage.prompt_tokens
+            self.total_completion_tokens += usage.completion_tokens
+            self.call_count += 1
+            self.usage_history.append(usage)
+
+            # Log token usage in structured format for analysis with color
+            model_short = usage.model_name.split("/")[-1] if "/" in usage.model_name else usage.model_name
+            token_log = f"TOKENS: {usage.total_tokens} | {usage.call_type} | {model_short} | P:{usage.prompt_tokens} C:{usage.completion_tokens}{'*' if usage.estimated else ''}"
+            logger.info(color_text(token_log, "magenta"))
+
+    def get_total_tokens(self) -> int:
+        """Get total tokens used."""
+        with self._lock:
+            return self.total_prompt_tokens + self.total_completion_tokens
+
+    def get_stats(self) -> Dict[str, Any]:
+        """Get comprehensive usage statistics."""
+        with self._lock:
+            session_duration = time.time() - self.session_start
+            return {
+                "session_duration_seconds": session_duration,
+                "total_calls": self.call_count,
+                "total_prompt_tokens": self.total_prompt_tokens,
+                "total_completion_tokens": self.total_completion_tokens,
+                "total_tokens": self.get_total_tokens(),
+                "usage_history": [
+                    {
+                        "prompt_tokens": usage.prompt_tokens,
+                        "completion_tokens": usage.completion_tokens,
+                        "total_tokens": usage.total_tokens,
+                        "model_name": usage.model_name,
+                        "timestamp": usage.timestamp,
+                        "call_type": usage.call_type,
+                        "estimated": usage.estimated,
+                    }
+                    for usage in self.usage_history
+                ],
+            }
+
+    def reset(self) -> None:
+        """Reset all usage statistics."""
+        with self._lock:
+            self.total_prompt_tokens = 0
+            self.total_completion_tokens = 0
+            self.session_start = time.time()
+            self.call_count = 0
+            self.usage_history.clear()
+            logger.info("Token usage tracker reset")
+
+
+# Global token tracker instance
+_global_tracker = TokenUsageTracker()
+
+
+def get_token_tracker() -> TokenUsageTracker:
+    """Get the global token usage tracker."""
+    return _global_tracker
+
+
+def record_token_usage(
+    prompt_tokens: int,
+    completion_tokens: int,
+    model_name: str,
+    call_type: str = "completion",
+    run_id: Optional[str] = None,
+    estimated: bool = False,
+) -> None:
+    """Record token usage in the global tracker."""
+    usage = TokenUsage(
+        prompt_tokens=prompt_tokens,
+        completion_tokens=completion_tokens,
+        total_tokens=prompt_tokens + completion_tokens,
+        model_name=model_name,
+        timestamp=datetime.now().isoformat(),
+        run_id=run_id,
+        call_type=call_type,
+        estimated=estimated,
+    )
+    _global_tracker.record_usage(usage)
+
+
+def extract_token_usage_from_openai_response(
+    response, model_name: str, call_type: str = "completion"
+) -> Optional[TokenUsage]:
+    """Extract token usage from OpenAI API response."""
+    try:
+        if hasattr(response, "usage") and response.usage:
+            usage = response.usage
+            return TokenUsage(
+                prompt_tokens=usage.prompt_tokens,
+                completion_tokens=usage.completion_tokens,
+                total_tokens=usage.total_tokens,
+                model_name=model_name,
+                timestamp=datetime.now().isoformat(),
+                call_type=call_type,
+                estimated=False,
+            )
+    except Exception as e:
+        logger.debug(f"Could not extract token usage from response: {e}")
+
+    return None
+
+
+def estimate_token_usage(
+    prompt_text: str, completion_text: str, model_name: str, call_type: str = "completion"
+) -> TokenUsage:
+    """Estimate token usage when actual usage is not available."""
+    # Simple estimation: ~4 characters per token for English text
+    prompt_tokens = max(1, len(prompt_text) // 4)
+    completion_tokens = max(1, len(completion_text) // 4)
+
+    return TokenUsage(
+        prompt_tokens=prompt_tokens,
+        completion_tokens=completion_tokens,
+        total_tokens=prompt_tokens + completion_tokens,
+        model_name=model_name,
+        timestamp=datetime.now().isoformat(),
+        call_type=call_type,
+        estimated=True,
+    )

noesium/core/utils/__init__.py

@@ -0,0 +1,10 @@
+"""
+Utility modules for noesium.core.
+"""
+
+from .logging import get_logger, setup_logging
+
+__all__ = [
+    "setup_logging",
+    "get_logger",
+]