nl2sql-engine 0.1.0__py3-none-any.whl

nl2sql/secrets/providers/hashi.py

@@ -0,0 +1,46 @@
+from typing import Optional
+import os
+import logging
+
+logger = logging.getLogger(__name__)
+
+class HashiCorpSecretProvider:
+    """Fetches secrets from HashiCorp Vault."""
+    
+    def __init__(self):
+        try:
+            import hvac
+            url = os.environ.get("VAULT_ADDR", "http://localhost:8200")
+            token = os.environ.get("VAULT_TOKEN")
+            
+            # Simple Token Auth for now, can extend to AppRole later
+            self.client = hvac.Client(url=url, token=token)
+            self._available = True
+        except ImportError:
+            self.client = None
+            self._available = False
+            logger.warning("HashiCorp Secret Provider initialized but 'hvac' missing. Install 'nl2sql-engine[hashicorp]'.")
+
+    def get_secret(self, key: str) -> Optional[str]:
+        if not self._available or not self.client:
+             raise ImportError("HashiCorp Secret Provider is not available.")
+             
+        try:
+            # key is assumed to be path/to/secret:key or just path (if single value?)
+            # Let's assume standard kv v2: "secret/data/my-app:password"
+            # Format: "mount/path:key"
+            
+            if ":" in key:
+                path, field = key.split(":", 1)
+            else:
+                path = key
+                field = "value" # Default field?
+
+            # Using KV V2
+            response = self.client.secrets.kv.v2.read_secret_version(path=path)
+            data = response['data']['data']
+            return data.get(field)
+            
+        except Exception as e:
+            logger.error(f"Failed to fetch secret '{key}' from Vault: {e}")
+            return None

nl2sql/services/callbacks/monitor.py

@@ -0,0 +1,84 @@
+from nl2sql.services.callbacks.node_handlers import NodeHandler
+from nl2sql.services.callbacks.token_handler import TokenHandler
+from langchain_core.callbacks import BaseCallbackHandler
+from nl2sql.services.callbacks.presenter import PresenterProtocol
+from typing import Dict, Any
+from langchain_core.outputs import LLMResult
+
+class PipelineMonitorCallback(BaseCallbackHandler):
+    """Callback handler for monitoring pipeline execution and auditing events.
+    
+    Integrates with OpenTelemetry for metrics and EventLogger for audit trails.
+    """
+    
+    def __init__(self, presenter: PresenterProtocol):
+        from nl2sql.common.settings import settings
+        from nl2sql.common.metrics import configure_metrics
+        
+        configure_metrics(
+            exporter_type=settings.observability_exporter,
+            otlp_endpoint=settings.otlp_endpoint
+        )
+        
+        self.node_handler = NodeHandler(presenter)
+        self.tokens = TokenHandler(self.node_handler.node_metrics)
+
+    def on_chain_start(self, serialized: Dict[str, Any], inputs: Dict[str, Any], **kwargs: Any) -> Any:
+        """Called when a chain starts."""
+        node_name = kwargs.get("metadata", {}).get("langgraph_node")
+        run_id = str(kwargs.get("run_id"))
+        parent_run_id = kwargs.get("parent_run_id")
+        parent_run_id = str(parent_run_id) if parent_run_id else None
+        self.node_handler.on_chain_start(run_id, parent_run_id, node_name, inputs)
+
+    def on_chain_end(self, outputs: Dict[str, Any], **kwargs: Any) -> Any:
+        """Called when a chain ends."""
+        run_id = str(kwargs.get("run_id"))
+        self.node_handler.on_chain_end(run_id)
+
+    def on_chain_error(self, error: BaseException, **kwargs: Any) -> Any:
+        """Called when a chain errors."""
+        run_id = str(kwargs.get("run_id"))
+        self.node_handler.on_chain_error(run_id, error)
+
+    def on_llm_end(self, response: LLMResult, **kwargs: Any) -> Any:
+        """Called when an LLM ends."""
+        from nl2sql.common.event_logger import event_logger
+        from nl2sql.common.logger import _trace_id_ctx, _tenant_id_ctx
+        
+        tags = kwargs.get("tags", [])
+        agent_name = next((t for t in tags if not t.startswith("seq:") and not t.startswith("langsmith:")), "unknown")
+        
+        text_output = ""
+        model_name = "unknown"
+        token_usage = {}
+        
+        if response.generations:
+            gen = response.generations[0][0]
+            text_output = gen.text
+            
+        if response.llm_output:
+            model_name = response.llm_output.get("model_name", "unknown")
+            token_usage = response.llm_output.get("token_usage", {})
+            
+        audit_payload = {
+            "agent": agent_name,
+            "model": model_name,
+            "response_snippet": text_output[:1000], 
+            "token_usage": token_usage
+        }
+        
+        event_logger.log_event(
+            event_type="llm_interaction", 
+            payload=audit_payload,
+            trace_id=_trace_id_ctx.get(),
+            tenant_id=_tenant_id_ctx.get()
+        )
+        
+        self.tokens.on_llm_end(response, agent_name=agent_name)
+
+    def get_status_tree(self):
+        self.node_handler.print_tree()
+
+    def get_performance_tree(self):
+        return self.node_handler.get_performance_tree()

nl2sql/services/callbacks/node_context.py

@@ -0,0 +1,7 @@
+import contextvars
+from typing import Optional
+
+current_node_run_id: contextvars.ContextVar[Optional[str]] = contextvars.ContextVar(
+    "current_node_run_id", default=None
+)
+

nl2sql/services/callbacks/node_handlers.py

@@ -0,0 +1,187 @@
+import time
+from typing import Dict, Any, List, Optional
+
+from nl2sql.common.context import current_datasource_id
+from nl2sql.common.metrics import LATENCY_LOG, node_duration_histogram
+from nl2sql.services.callbacks.presenter import PresenterProtocol
+from nl2sql.services.callbacks.node_context import current_node_run_id
+from nl2sql.services.callbacks.node_metrics import NodeMetrics
+
+
+class NodeHandler:
+    """Handles node execution lifecycle events and metrics recording."""
+
+    def __init__(self, presenter: PresenterProtocol):
+        self.presenter = presenter
+
+        self.run_start: Dict[str, float] = {}
+        self.run_to_node: Dict[str, str] = {}
+        self.run_parent: Dict[str, Optional[str]] = {}
+
+        self.primary_run_id: Dict[str, str] = {}
+
+        self.tree: Dict[str, List[str]] = {"root": []}
+        self.node_metrics: Dict[str, NodeMetrics] = {}
+
+        self.node_active_count: Dict[str, int] = {}
+        self.run_ctx_tokens: Dict[str, Any] = {}
+        self.ds_ctx_tokens: Dict[str, Any] = {}
+
+    def on_chain_start(
+        self,
+        run_id: str,
+        parent_run_id: Optional[str],
+        node_name: Optional[str],
+        inputs: Dict[str, Any],
+    ):
+        if not node_name:
+            return
+
+        now = time.perf_counter()
+
+        raw_node = node_name
+        display_node = raw_node
+
+        if raw_node == "execution_branch":
+            if isinstance(inputs, dict):
+                q = inputs.get("user_query")
+                if not q:
+                    sub_query = inputs.get("sub_query")
+                    if isinstance(sub_query, dict):
+                        q = sub_query.get("intent")
+                    elif hasattr(sub_query, "intent"):
+                        q = sub_query.intent
+                if q:
+                    display_node = f"execution_branch ({q})"
+
+        self.run_start[run_id] = now
+        self.run_to_node[run_id] = display_node
+        self.run_parent[run_id] = parent_run_id
+
+        if display_node not in self.primary_run_id:
+            self.primary_run_id[display_node] = run_id
+
+            parent_primary = "root"
+            if parent_run_id and parent_run_id in self.run_to_node:
+                parent_node = self.run_to_node[parent_run_id]
+                parent_primary = self.primary_run_id.get(parent_node, "root")
+
+            self.tree.setdefault(parent_primary, []).append(run_id)
+            self.tree.setdefault(run_id, [])
+
+            self.node_metrics[run_id] = NodeMetrics(
+                start_time=now,
+                end_time=now,
+                datasource_id=current_datasource_id.get(),
+            )
+
+        primary_id = self.primary_run_id[display_node]
+        metrics = self.node_metrics[primary_id]
+        metrics.start_time = min(metrics.start_time, now)
+
+        ds_id = None
+        if isinstance(inputs, dict):
+            ds_id = inputs.get("selected_datasource_id") or inputs.get("datasource_id")
+            if not ds_id:
+                sub_query = inputs.get("sub_query")
+                if isinstance(sub_query, dict):
+                    ds_id = sub_query.get("datasource_id")
+                elif hasattr(sub_query, "datasource_id"):
+                    ds_id = sub_query.datasource_id
+
+        if ds_id:
+            tok = current_datasource_id.set(ds_id)
+            self.ds_ctx_tokens[run_id] = tok
+
+        tok = current_node_run_id.set(primary_id)
+        self.run_ctx_tokens[run_id] = tok
+
+        self.node_active_count[display_node] = (
+            self.node_active_count.get(display_node, 0) + 1
+        )
+
+        if self.node_active_count[display_node] == 1:
+            self.presenter.update_interactive_status(
+                f"[bold blue]{display_node}[/bold blue] Working..."
+            )
+
+    def on_chain_end(self, run_id: str):
+        node = self.run_to_node.get(run_id)
+        start = self.run_start.get(run_id)
+
+        if not node or start is None:
+            return
+
+        end = time.perf_counter()
+        primary_id = self.primary_run_id[node]
+        metrics = self.node_metrics[primary_id]
+
+        metrics.end_time = max(metrics.end_time, end)
+        metrics.duration = metrics.end_time - metrics.start_time
+
+
+        LATENCY_LOG.append(
+            {
+                "node": node,
+                "duration": end - start,
+                "datasource_id": current_datasource_id.get(),
+                "run_id": run_id,
+                "parent_run_id": self.run_parent.get(run_id),
+            }
+        )
+        
+        node_duration_histogram.record(
+            end - start,
+            attributes={
+                "node": node,
+                "datasource_id": str(current_datasource_id.get() or "none"),
+            }
+        )
+
+        self.node_active_count[node] -= 1
+        if self.node_active_count[node] == 0:
+            tok_str = f" | {metrics.total_tokens} tok" if metrics.total_tokens else ""
+            self.presenter.print_success(
+                f"{node} Completed ({metrics.duration:.2f}s{tok_str})"
+            )
+            self.presenter.update_interactive_status("Thinking...")
+
+        tok = self.run_ctx_tokens.pop(run_id, None)
+        if tok:
+            current_node_run_id.reset(tok)
+
+        ds_tok = self.ds_ctx_tokens.pop(run_id, None)
+        if ds_tok:
+            current_datasource_id.reset(ds_tok)
+
+    def on_chain_error(self, run_id: str, error: BaseException):
+        node = self.run_to_node.get(run_id)
+        start = self.run_start.get(run_id)
+        end = time.perf_counter()
+
+        if node:
+            primary_id = self.primary_run_id[node]
+            metrics = self.node_metrics[primary_id]
+            metrics.end_time = max(metrics.end_time, end)
+            metrics.duration = metrics.end_time - metrics.start_time
+            metrics.error = str(error)
+
+            self.presenter.print_error(
+                f"{node} Failed: {error}"
+            )
+            self.presenter.update_interactive_status("Error encountered...")
+
+        tok = self.run_ctx_tokens.pop(run_id, None)
+        if tok:
+            current_node_run_id.reset(tok)
+
+        ds_tok = self.ds_ctx_tokens.pop(run_id, None)
+        if ds_tok:
+            current_datasource_id.reset(ds_tok)
+
+    def get_performance_tree(self):
+        return (
+            self.tree,
+            self.node_metrics,
+            {rid: self.run_to_node[rid] for rid in self.node_metrics},
+        )

nl2sql/services/callbacks/node_metrics.py

@@ -0,0 +1,14 @@
+from dataclasses import dataclass
+from typing import Optional
+
+
+@dataclass
+class NodeMetrics:
+    start_time: float = float("inf")
+    end_time: float = 0.0
+    duration: float = 0.0
+    total_tokens: int = 0
+    prompt_tokens: int = 0
+    completion_tokens: int = 0
+    error: Optional[str] = None
+    datasource_id: Optional[str] = None

nl2sql/services/callbacks/presenter.py

@@ -0,0 +1,12 @@
+from typing import Protocol
+
+
+class PresenterProtocol(Protocol):
+    def update_interactive_status(self, message: str) -> None:
+        ...
+
+    def print_success(self, message: str) -> None:
+        ...
+
+    def print_error(self, message: str) -> None:
+        ...

nl2sql/services/callbacks/token_handler.py

@@ -0,0 +1,56 @@
+from langchain_core.outputs import LLMResult
+from nl2sql.common.metrics import TOKEN_LOG, token_usage_counter
+from nl2sql.common.context import current_datasource_id
+from nl2sql.services.callbacks.node_context import current_node_run_id
+from nl2sql.services.callbacks.node_metrics import NodeMetrics
+
+
+class TokenHandler:
+    """Handles token usage tracking and metrics."""
+
+    def __init__(self, node_metrics: dict[str, NodeMetrics]):
+        self.node_metrics = node_metrics
+
+    def on_llm_end(self, response: LLMResult, agent_name: str = "unknown", model_name: str = "unknown"):
+        """Records token usage from LLM response."""
+        usage = None
+        if response and response.llm_output:
+            usage = response.llm_output.get("token_usage") or response.llm_output.get("usage")
+
+        if not usage:
+            return
+
+        p = usage.get("prompt_tokens") or usage.get("input_tokens", 0)
+        c = usage.get("completion_tokens") or usage.get("output_tokens", 0)
+        t = usage.get("total_tokens") or (p + c)
+
+        run_id = current_node_run_id.get()
+
+
+        TOKEN_LOG.append(
+            {
+                "agent": agent_name,
+                "model": model_name,
+                "datasource_id": current_datasource_id.get(),
+                "prompt_tokens": p,
+                "completion_tokens": c,
+                "total_tokens": t,
+                "run_id": run_id,
+            }
+        )
+
+        token_usage_counter.add(
+            t,
+            attributes={
+                "agent": agent_name,
+                "model": model_name,
+                "datasource_id": str(current_datasource_id.get() or "none"),
+                "type": "total"
+            }
+        )
+
+        if run_id and run_id in self.node_metrics:
+            m = self.node_metrics[run_id]
+            m.prompt_tokens += int(p)
+            m.completion_tokens += int(c)
+            m.total_tokens += int(t)