qaiss-sdk 0.1.0__py3-none-any.whl

qaiss/__init__.py

@@ -0,0 +1,70 @@
+"""
+QAISS SDK: Production-grade Observability, Telemetry & Guardrails Client for AI Agents.
+Collects and structures ground-truth telemetry for the QAISS Hallucination & Safety Evaluation Engine.
+"""
+
+from typing import Optional, Any
+from .client import QaissClient
+from .context import set_context, get_current_security_context, clear_context, get_next_step_index
+from .models import QaissEventModel, SecurityContext, LLMMetadata, EventType
+from .redaction import scrub_payload, redact_string, is_valid_luhn
+from .truncation import truncate_payload
+from .core import QaissRequestEngine
+from .asgi import QaissASGIMiddleware
+from .wsgi import QaissWSGIMiddleware
+from .django_middleware import QaissDjangoMiddleware
+from .middleware import QaissMiddleware, QaissFastAPIMiddleware
+from .instrumentation import wrap_tools, auto_instrument
+
+__version__ = "0.1.0"
+_global_client: Optional[QaissClient] = None
+
+def init(
+    api_key: Optional[str] = None,
+    backend_url: Optional[str] = None,
+    environment: Optional[str] = None,
+    debug: bool = False,
+    **kwargs: Any
+) -> QaissClient:
+    """
+    Sentry-style single-line automatic initialization.
+    Automatically patches all AI libraries (OpenAI, Anthropic, Requests, HTTPX)
+    and Agent Framework constructors (Agno, LangChain, CrewAI, LlamaIndex) process-wide.
+    """
+    global _global_client
+    _global_client = QaissClient(
+        api_key=api_key,
+        backend_url=backend_url,
+        environment=environment,
+        debug=debug,
+        auto_track_libraries=True,
+        **kwargs
+    )
+    return _global_client
+
+__all__ = [
+    "init",
+    "QaissClient",
+    "QaissRequestEngine",
+    "QaissASGIMiddleware",
+    "QaissWSGIMiddleware",
+    "QaissDjangoMiddleware",
+    "QaissMiddleware",
+    "QaissFastAPIMiddleware",
+    "wrap_tools",
+    "auto_instrument",
+    "set_context",
+    "get_current_security_context",
+    "clear_context",
+    "get_next_step_index",
+    "QaissEventModel",
+    "SecurityContext",
+    "LLMMetadata",
+    "EventType",
+    "scrub_payload",
+    "redact_string",
+    "is_valid_luhn",
+    "truncate_payload"
+]
+
+__version__ = "0.1.0"

qaiss/asgi.py

@@ -0,0 +1,117 @@
+"""
+QAISS Universal ASGI Middleware. Speaks raw ASGI 3.0 — no Starlette/FastAPI import.
+Works with FastAPI, Starlette, Litestar, Quart, BlackSheep, Django (ASGI mode), and
+any other ASGI-compliant framework, current or future, with zero new SDK code.
+
+IMPORTANT: unlike Starlette's BaseHTTPMiddleware (which must fully buffer
+response.body_iterator before anything reaches the client), this middleware forwards
+every send() call to the real transport immediately. Streaming/SSE responses (e.g.
+token-by-token chat output) are NOT delayed or buffered for the client — only our own
+internal copy is buffered, capped by max_body_bytes.
+"""
+import logging
+from typing import Any, List, Optional
+
+from .core import QaissRequestEngine
+
+logger = logging.getLogger("qaiss.asgi")
+
+
+class QaissASGIMiddleware:
+    def __init__(
+        self,
+        app,
+        qaiss_client: Any,
+        include_paths: Optional[List[str]] = None,
+        max_body_bytes: int = 1_000_000,
+        **engine_kwargs,
+    ):
+        self.app = app
+        self.engine = QaissRequestEngine(qaiss_client, **engine_kwargs)
+        self.include_paths = set(include_paths) if include_paths else None
+        self.max_body_bytes = max_body_bytes
+
+    def _should_track(self, path: str) -> bool:
+        if not self.include_paths:
+            return True
+        return any(path.startswith(p) for p in self.include_paths)
+
+    async def __call__(self, scope, receive, send):
+        # Non-HTTP scopes (websocket, lifespan) pass straight through untouched.
+        if scope["type"] != "http" or not self._should_track(scope.get("path", "")):
+            await self.app(scope, receive, send)
+            return
+
+        headers = {k.decode("latin-1"): v.decode("latin-1") for k, v in scope.get("headers", [])}
+        method = scope.get("method", "GET")
+        path = scope.get("path", "")
+
+        content_length_header = headers.get("content-length")
+        content_length = int(content_length_header) if content_length_header and content_length_header.isdigit() else None
+
+        # If request is explicitly larger than max_body_bytes, do not pre-read or buffer.
+        # Pass receive directly to downstream app so streaming uploads work with 0 memory overhead.
+        if content_length is not None and content_length > self.max_body_bytes:
+            state = self.engine.before_request(method, path, headers, None)
+            app_receive = receive
+        else:
+            buffered_messages = []
+            body_chunks = []
+            total_bytes = 0
+            more_body = True
+            capture_request_body = True
+
+            while more_body:
+                message = await receive()
+                buffered_messages.append(message)
+                chunk = message.get("body", b"")
+                total_bytes += len(chunk)
+                if total_bytes <= self.max_body_bytes and capture_request_body:
+                    body_chunks.append(chunk)
+                else:
+                    capture_request_body = False
+                more_body = message.get("more_body", False)
+
+            body_bytes = b"".join(body_chunks) if capture_request_body else None
+            state = self.engine.before_request(method, path, headers, body_bytes)
+
+            msg_idx = 0
+
+            async def replay_receive():
+                nonlocal msg_idx
+                if msg_idx < len(buffered_messages):
+                    msg = buffered_messages[msg_idx]
+                    msg_idx += 1
+                    return msg
+                return await receive()
+
+            app_receive = replay_receive
+
+        response_status = {"code": 200}
+        response_headers_list = []
+        response_body_chunks = []
+        bytes_seen = {"total": 0}
+        capture_response = {"on": True}
+
+        async def send_wrapper(message):
+            if message["type"] == "http.response.start":
+                response_status["code"] = message["status"]
+                response_headers_list.extend(
+                    (k.decode("latin-1"), v.decode("latin-1")) for k, v in message.get("headers", [])
+                )
+            elif message["type"] == "http.response.body":
+                chunk = message.get("body", b"")
+                bytes_seen["total"] += len(chunk)
+                if capture_response["on"]:
+                    if bytes_seen["total"] > self.max_body_bytes:
+                        capture_response["on"] = False  # stop buffering our copy; client is unaffected
+                    else:
+                        response_body_chunks.append(chunk)
+            await send(message)  # ALWAYS forward immediately — this preserves streaming
+
+        try:
+            await self.app(scope, app_receive, send_wrapper)
+        finally:
+            resp_headers = dict(response_headers_list)
+            resp_body = b"".join(response_body_chunks) if capture_response["on"] else None
+            self.engine.after_request(state, response_status["code"], resp_headers, resp_body)

qaiss/client.py

@@ -0,0 +1,283 @@
+import os
+import random
+import uuid
+import json
+import logging
+from typing import Optional, List, Dict, Any, Set, Union
+from contextlib import contextmanager
+
+from opentelemetry import trace
+from opentelemetry.sdk.trace import TracerProvider
+from opentelemetry.trace import ProxyTracerProvider
+from opentelemetry.trace.propagation.tracecontext import TraceContextTextMapPropagator
+
+from .transport import BackgroundTransport
+from .decorator import create_track_tool_decorator, create_track_agent_decorator, format_utc_timestamp
+from .models import QaissEventModel, SecurityContext, LLMMetadata, EventType
+from .redaction import scrub_payload
+from .truncation import truncate_payload
+from .context import set_context, get_current_security_context, clear_context, get_next_step_index
+from .instrumentation import wrap_tools, auto_instrument
+
+logger = logging.getLogger("qaiss.client")
+_propagator = TraceContextTextMapPropagator()
+
+class QaissClient:
+    """
+    Production-grade Observability, Telemetry & Guardrails Client for AI Agents.
+
+    Dual-Purpose Instrumentation:
+    1. Global Process-Wide Auto-Tracking (No arguments):
+       Call `client.wrap_tools()` to automatically patch installed libraries (OpenAI, Anthropic, Requests, HTTPX).
+       All outgoing calls across your entire app are tracked automatically with zero further code.
+    
+    2. Explicit Tool Wrapping (With arguments):
+       Call `client.wrap_tools(tools)` to instrument specific tool lists, LangChain/LlamaIndex/CrewAI/Agno tools,
+       or custom class instances.
+
+    3. Manual Decorators:
+       Use `@client.track_tool(...)` for proprietary DB or internal API functions.
+    """
+    def __init__(
+        self,
+        api_key: Optional[str] = None,
+        backend_url: Optional[str] = None,
+        environment: Optional[str] = None,
+        enabled: bool = True,
+        debug: bool = False,
+        sample_rate: float = 1.0,
+        tracked_users: Optional[List[str]] = None,
+        include_tools: Optional[List[str]] = None,
+        exclude_tools: Optional[List[str]] = None,
+        client_id: Optional[str] = None,
+        auto_track_libraries: bool = True,
+        track_db: bool = False
+    ):
+        self.api_key = api_key or os.environ.get("QAISS_API_KEY", "")
+        self.backend_url = backend_url or os.environ.get("QAISS_BACKEND_URL", "https://api.qaiss.com/v1/events")
+        self.environment = environment or os.environ.get("QAISS_ENVIRONMENT", "production")
+        self.enabled = enabled
+        self.debug = debug
+        self.sample_rate = max(0.0, min(1.0, sample_rate))
+        self.client_id = client_id or f"client_{uuid.uuid4().hex[:8]}"
+        self.auto_track_libraries = auto_track_libraries
+        self.track_db = track_db
+        self._auto_tracked_libraries: Dict[str, bool] = {}
+        
+        self.tracked_users: Optional[Set[str]] = {str(u) for u in tracked_users} if tracked_users else None
+        self.include_tools: Optional[Set[str]] = set(include_tools) if include_tools else None
+        self.exclude_tools: Optional[Set[str]] = set(exclude_tools) if exclude_tools else None
+
+        provider = trace.get_tracer_provider()
+        if isinstance(provider, ProxyTracerProvider) or provider is None:
+            self._tracer_provider = TracerProvider()
+            self._tracer = self._tracer_provider.get_tracer(f"qaiss-sdk-{self.client_id}")
+        else:
+            self._tracer = trace.get_tracer(f"qaiss-sdk-{self.client_id}")
+
+        self._transport = BackgroundTransport(backend_url=self.backend_url, api_key=self.api_key)
+        
+        self.track_tool = create_track_tool_decorator(self)
+        self.track_agent = create_track_agent_decorator(self)
+
+        # Opt-in automatic process-wide library patching at construction time
+        if self.auto_track_libraries and self.enabled:
+            self._auto_tracked_libraries = self.wrap_tools()
+
+    def wrap_tools(
+        self,
+        tools: Any = None,
+        track_db: Optional[bool] = None
+    ) -> Union[Dict[str, bool], List[Any], Dict[str, Any], Any]:
+        """
+        Dual-Purpose Tool & Library Instrumentor.
+
+        Modes:
+        - `wrap_tools()` (no args): Auto-patches all installed libraries (OpenAI, Anthropic, Gemini, Requests, HTTPX, Agent frameworks)
+        - `wrap_tools(tools)`: Wraps functions, toolkits, LangChain/CrewAI tools, or custom class methods.
+        """
+        use_track_db = self.track_db if track_db is None else track_db
+        return wrap_tools(tools=tools, client=self, track_db=use_track_db)
+
+    def auto_instrument(self, frameworks: Optional[List[str]] = None) -> Dict[str, bool]:
+        """
+        [DEPRECATED] Use client.wrap_tools() with no arguments instead.
+        Auto-instruments installed AI and HTTP libraries process-wide.
+        """
+        logger.debug("QAISS: auto_instrument() is deprecated and will be removed in a future release. Use wrap_tools() with no arguments instead.")
+        return self.wrap_tools()
+
+    def get_tracer(self):
+        return self._tracer
+
+    def set_context(
+        self,
+        authenticated_user_id: str,
+        resource_owner_id: Optional[str] = None,
+        tenant_id: Optional[str] = None,
+        scopes_provided: Optional[List[str]] = None
+    ) -> None:
+        set_context(
+            authenticated_user_id=authenticated_user_id,
+            resource_owner_id=resource_owner_id,
+            tenant_id=tenant_id,
+            scopes_provided=scopes_provided,
+            client_id=self.client_id
+        )
+
+    def get_context(self) -> Optional[Dict[str, Any]]:
+        return get_current_security_context(client_id=self.client_id)
+
+    def clear_context(self) -> None:
+        clear_context(client_id=self.client_id)
+
+    def get_next_step_index(self) -> int:
+        return get_next_step_index(client_id=self.client_id)
+
+    def get_active_traceparent(self) -> str:
+        carrier = {}
+        _propagator.inject(carrier)
+        return carrier.get("traceparent", "")
+
+    def should_record(self, tool_name: str, security_context: Optional[Dict[str, Any]] = None) -> bool:
+        if not self.enabled:
+            return False
+
+        if self.exclude_tools and tool_name in self.exclude_tools:
+            return False
+
+        if self.include_tools is not None and tool_name not in self.include_tools:
+            return False
+
+        if self.tracked_users is not None:
+            auth_user = (security_context or {}).get("authenticated_user_id")
+            if not auth_user or str(auth_user) not in self.tracked_users:
+                return False
+
+        if self.sample_rate < 1.0:
+            if random.random() > self.sample_rate:
+                return False
+
+        return True
+
+    def _print_debug_event(self, event_dict: Dict[str, Any]):
+        """Prints a formatted flight recorder box directly in terminal for the developer."""
+        step = event_dict.get("step_index", 1)
+        event_type = event_dict.get("event_type", "event")
+        name = event_dict.get("name", "unknown")
+        status = event_dict.get("status", "success")
+        duration = event_dict.get("duration_ms", 0)
+        traceparent = event_dict.get("traceparent", "N/A")
+        redactions = event_dict.get("redacted_fields")
+        user = (event_dict.get("security_context") or {}).get("authenticated_user_id", "N/A")
+        status_icon = "[OK]" if status == "success" else "[ERR]"
+        sep = "-" * 80
+
+        try:
+            lines = [
+                "",
+                sep,
+                f"[QAISS] STEP {step} | {str(event_type).upper()} : '{name}' {status_icon}",
+                sep,
+                f"  User ID  : {user:<24} Duration : {duration} ms",
+                f"  Trace ID : {traceparent}",
+            ]
+            if redactions:
+                lines.append(f"  Scrubbed : {redactions}")
+            if event_dict.get("input"):
+                preview = json.dumps(event_dict["input"], indent=2, ensure_ascii=True)
+                if len(preview) > 400:
+                    preview = preview[:400] + "\n  ...[truncated]"
+                lines.append(f"  Input:\n{preview}")
+            if event_dict.get("output"):
+                preview = json.dumps(event_dict["output"], indent=2, ensure_ascii=True)
+                if len(preview) > 400:
+                    preview = preview[:400] + "\n  ...[truncated]"
+                lines.append(f"  Output:\n{preview}")
+            lines.append(sep)
+            print("\n".join(lines))
+        except Exception as e:
+            logger.debug(f"QAISS: debug print failed: {e}")
+
+    def track_response(
+        self,
+        final_message: str,
+        user_prompt: Optional[str] = None,
+        security_context: Optional[Dict[str, Any]] = None,
+        llm_metadata: Optional[Dict[str, Any]] = None,
+        workflow_id: Optional[str] = None
+    ):
+        """Manually tracks the user-visible final chatbot response."""
+        try:
+            effective_sec_ctx = security_context or self.get_context()
+
+            if not self.should_record("chatbot_response", effective_sec_ctx):
+                return
+
+            step_index = self.get_next_step_index()
+
+            # PII scrubbing
+            clean_msg, msg_redactions = scrub_payload(final_message, "output.final_response")
+            clean_prompt, prompt_redactions = scrub_payload(user_prompt, "input.user_prompt") if user_prompt else (None, [])
+            all_redactions = (msg_redactions + prompt_redactions) or None
+
+            # Truncation
+            final_msg_trunc, is_msg_trunc = truncate_payload(clean_msg)
+            final_prompt_trunc, is_prompt_trunc = truncate_payload(clean_prompt) if clean_prompt else (None, False)
+
+            sec_ctx = SecurityContext(**effective_sec_ctx) if effective_sec_ctx else None
+            llm_meta = LLMMetadata(**llm_metadata) if llm_metadata else None
+
+            now_ts = format_utc_timestamp()
+
+            event = QaissEventModel(
+                event_id=f"evt_{uuid.uuid4().hex[:12]}",
+                traceparent=self.get_active_traceparent(),
+                workflow_id=workflow_id,
+                step_index=step_index,
+                event_type="final_response",
+                name="chatbot_response",
+                status="success",
+                start_timestamp=now_ts,
+                end_timestamp=now_ts,
+                duration_ms=0,
+                input={"user_prompt": final_prompt_trunc} if user_prompt else None,
+                output={"final_response": final_msg_trunc},
+                is_truncated=(is_msg_trunc or is_prompt_trunc),
+                redacted_fields=all_redactions,
+                security_context=sec_ctx,
+                llm_metadata=llm_meta,
+                tags={"environment": self.environment}
+            )
+            self.send_event(event.to_wire_dict())
+        except Exception as e:
+            logger.debug(f"QAISS: track_response failed safely: {e}")
+
+    @contextmanager
+    def trace_session(self, session_name: str = "agent_session"):
+        with self._tracer.start_as_current_span(session_name) as span:
+            yield span
+
+    def send_event(self, event_dict: Dict[str, Any]):
+        try:
+            if self.debug:
+                self._print_debug_event(event_dict)
+            self._transport.enqueue(event_dict)
+        except Exception as e:
+            logger.debug(f"QAISS: send_event failed safely: {e}")
+
+    def flush(self, timeout: float = 2.0):
+        try:
+            if hasattr(self._transport, "flush"):
+                self._transport.flush(timeout=timeout)
+        except Exception:
+            pass
+
+    def close(self, timeout: float = 2.0):
+        try:
+            if hasattr(self._transport, "flush_and_close"):
+                self._transport.flush_and_close()
+            elif hasattr(self._transport, "close"):
+                self._transport.close()
+        except Exception:
+            pass

qaiss/context.py

@@ -0,0 +1,80 @@
+from contextvars import ContextVar
+from typing import Optional, Dict, Any, List
+
+# Multi-client, thread-safe, async-isolated request context variables
+_contexts_by_client: ContextVar[Dict[str, Optional[Dict[str, Any]]]] = ContextVar("qaiss_contexts_by_client", default={})
+_step_counters_by_client: ContextVar[Dict[str, int]] = ContextVar("qaiss_step_counters_by_client", default={})
+
+def set_context(
+    authenticated_user_id: str,
+    resource_owner_id: Optional[str] = None,
+    tenant_id: Optional[str] = None,
+    scopes_provided: Optional[List[str]] = None,
+    client_id: str = "default"
+) -> None:
+    """
+    Sets the request-scoped security context and resets the step counter for a specific client instance.
+    Guarantees full isolation across multiple QaissClient instances (multi-tenant or multi-service).
+    """
+    # Copy dict for ContextVar immutability across async coroutines
+    curr_counters = dict(_step_counters_by_client.get())
+    curr_counters[client_id] = 0
+    _step_counters_by_client.set(curr_counters)
+
+    curr_contexts = dict(_contexts_by_client.get())
+    curr_contexts[client_id] = {
+        "authenticated_user_id": str(authenticated_user_id) if authenticated_user_id is not None else None,
+        "resource_owner_id": str(resource_owner_id) if resource_owner_id is not None else (str(authenticated_user_id) if authenticated_user_id is not None else None),
+        "tenant_id": str(tenant_id) if tenant_id is not None else None,
+        "scopes_provided": scopes_provided
+    }
+    _contexts_by_client.set(curr_contexts)
+
+def get_current_security_context(client_id: str = "default") -> Optional[Dict[str, Any]]:
+    """Retrieves the request-scoped security context for a specific client instance."""
+    contexts = _contexts_by_client.get()
+    if not contexts:
+        return None
+    if client_id in contexts and contexts[client_id]:
+        return contexts[client_id]
+    if "default" in contexts and contexts["default"]:
+        return contexts["default"]
+    
+    # Fallback to single active context ONLY for un-scoped/default lookups (global library patches)
+    if client_id == "default":
+        non_empty = [c for c in contexts.values() if c]
+        if len(non_empty) == 1:
+            return non_empty[0]
+
+    return None
+
+def get_next_step_index(client_id: str = "default") -> int:
+    """Increments and returns the sequential step index for the given client instance."""
+    curr_counters = dict(_step_counters_by_client.get())
+    
+    # If the requested client_id is default and not yet in counters, but exactly 1 client context is active
+    target_id = client_id
+    if target_id == "default" and target_id not in curr_counters and len(curr_counters) == 1:
+        target_id = next(iter(curr_counters.keys()))
+
+    current = curr_counters.get(target_id, 0)
+    next_step = current + 1
+    curr_counters[target_id] = next_step
+    if client_id != target_id:
+        curr_counters[client_id] = next_step
+    _step_counters_by_client.set(curr_counters)
+    return next_step
+
+def clear_context(client_id: Optional[str] = None) -> None:
+    """Clears the security context and resets the step counter for a client (or all clients)."""
+    if client_id:
+        curr_contexts = dict(_contexts_by_client.get())
+        curr_contexts.pop(client_id, None)
+        _contexts_by_client.set(curr_contexts)
+
+        curr_counters = dict(_step_counters_by_client.get())
+        curr_counters.pop(client_id, None)
+        _step_counters_by_client.set(curr_counters)
+    else:
+        _contexts_by_client.set({})
+        _step_counters_by_client.set({})