go-otel/cmd/needle-sidecar/server.py

#!/usr/bin/env python3
"""
Needle 3 AI Telemetry Triage Sidecar Service
Accepts OTel incident summaries and uses Cactus-Compute/needle3 to classify and extract actionable insights.
"""

import os
import sys
import json
import time
import re
import hashlib
from http.server import ThreadingHTTPServer, BaseHTTPRequestHandler
import needle

PORT = int(os.environ.get("NEEDLE_PORT", 5005))
HOST = os.environ.get("NEEDLE_HOST", "0.0.0.0")

# --- Cause-level triage tools ---
# These deliberately do not restate what the event already carries (status
# code, is_error, severity). Each one is a judgement about *why* it happened
# or *who should care*, which is what an alert reader needs.
@needle.tool
def report_scanner_probe(service: str, path: str, status_code: int):
    """Automated scan or exploit crawler requesting sensitive or nonexistent paths such as .env, wp-admin, .git, phpmyadmin. Applies whatever status the probe received, including 500 or 502."""
    return {"action": "scanner_probe", "service": service, "path": path, "status_code": status_code}

@needle.tool
def report_dependency_unreachable(service: str, dependency: str, error_message: str):
    """A downstream or upstream service is down, refusing connections, closed the stream (EOF), or returned upstream unavailable."""
    return {"action": "dependency_unreachable", "service": service, "dependency": dependency, "error": error_message}

@needle.tool
def report_client_credential_issue(service: str, route: str, reason: str):
    """A caller lacks a valid or unexpired token, for example an unauthenticated client polling an API or a registry pull without login. The caller's fault, not a fault in this service."""
    return {"action": "client_credential_issue", "service": service, "route": route, "reason": reason}

@needle.tool
def report_stale_reference(service: str, path: str):
    """Harmless missing resource caused by an old link, removed asset, deleted commit, or dev tooling request."""
    return {"action": "stale_reference", "service": service, "path": path}

@needle.tool
def report_service_defect(service: str, error_message: str):
    """A bug in the service's own code: panic, nil dereference, unhandled exception, index out of range, deadlock, or a 500 from a real application route."""
    return {"action": "service_defect", "service": service, "error": error_message}

@needle.tool
def report_capacity_saturation(service: str, resource: str, error_message: str):
    """Resource exhaustion or overload: timeouts, connection pool exhausted, memory limit exceeded, OOM, queue full."""
    return {"action": "capacity_saturation", "service": service, "resource": resource, "error": error_message}

@needle.tool
def report_expected_lifecycle(service: str, what: str):
    """Normal behaviour that is not a problem: health and readiness checks, keepalive pings, retries, cancellations, client disconnects, redirects, successful requests, routine info logs, web vitals."""
    return {"action": "expected_lifecycle", "service": service, "what": what}

@needle.tool
def report_needs_review(service: str, summary: str):
    """Fallback when none of the other categories clearly applies."""
    return {"action": "needs_review", "service": service, "summary": summary}

# --- Redaction: nothing PII-bearing should reach the model or the cache ---
_REDACTIONS = [
    (re.compile(r"[\w.+-]+@[\w-]+\.[\w.-]+"), "<email>"),
    (re.compile(r"\b(?:\d{1,3}\.){3}\d{1,3}\b"), "<ip>"),
    (re.compile(r"\b[0-9a-fA-F]{8}-(?:[0-9a-fA-F]{4}-){3}[0-9a-fA-F]{12}\b"), "<uuid>"),
    (re.compile(r"\beyJ[\w-]+\.[\w-]+\.[\w-]+\b"), "<jwt>"),
    (re.compile(r"\b[0-9a-fA-F]{16,}\b"), "<hex>"),
    (re.compile(r"([?&][\w.-]+=)[^&\s|]+"), r"\1<v>"),
]

def redact(text: str) -> str:
    for rx, sub in _REDACTIONS:
        text = rx.sub(sub, text)
    return text[:1000]

# Cache keyed on the redacted text: repeated signatures cost nothing.
CACHE_MAX = int(os.environ.get("NEEDLE_CACHE_MAX", 2000))
_cache = {}
_stats = {"hits": 0, "misses": 0}

# Initialize Needle Agent once
print(f"Loading Needle 3 model (Cactus-Compute/needle3)...")
t0 = time.perf_counter()
agent = needle.Needle(tools=[
    report_scanner_probe,
    report_dependency_unreachable,
    report_client_credential_issue,
    report_stale_reference,
    report_service_defect,
    report_capacity_saturation,
    report_expected_lifecycle,
    report_needs_review,
])
print(f"Needle 3 loaded in {(time.perf_counter() - t0):.2f}s")


class NeedleHandler(BaseHTTPRequestHandler):
    def _send_cors(self):
        self.send_header("Access-Control-Allow-Origin", "*")
        self.send_header("Access-Control-Allow-Methods", "GET, POST, OPTIONS")
        self.send_header("Access-Control-Allow-Headers", "Content-Type, Authorization, X-Otel-Agent-Token")

    def do_OPTIONS(self):
        self.send_response(204)
        self._send_cors()
        self.end_headers()

    def do_GET(self):
        if self.path == "/health" or self.path == "/api/needle/status":
            self.send_response(200)
            self._send_cors()
            self.send_header("Content-Type", "application/json")
            self.end_headers()
            resp = {
                "status": "ok",
                "available": True,
                "model": "Cactus-Compute/needle3",
                "cache": {**_stats, "size": len(_cache)},
                "tools": [
                    "scanner_probe", "dependency_unreachable", "client_credential_issue", "stale_reference",
                    "service_defect", "capacity_saturation", "expected_lifecycle", "needs_review"
                ]
            }
            self.wfile.write(json.dumps(resp).encode("utf-8"))
        else:
            self.send_response(404)
            self._send_cors()
            self.end_headers()

    def do_POST(self):
        if self.path in ("/classify", "/api/needle/classify"):
            content_len = int(self.headers.get("Content-Length", 0))
            body = self.rfile.read(content_len)
            try:
                data = json.loads(body) if body else {}
            except Exception as e:
                self.send_response(400)
                self._send_cors()
                self.send_header("Content-Type", "application/json")
                self.end_headers()
                self.wfile.write(json.dumps({"error": f"Invalid JSON: {e}"}).encode("utf-8"))
                return

            text = data.get("text", "").strip()
            # If text is not provided, build from structured fields:
            if not text:
                parts = []
                if data.get("service"):
                    parts.append(f"Service: {data['service']}")
                if data.get("kind"):
                    parts.append(f"Kind: {data['kind']}")
                if data.get("title"):
                    parts.append(f"Title: {data['title']}")
                if data.get("message"):
                    parts.append(f"Message: {data['message']}")
                if data.get("url"):
                    parts.append(f"URL: {data['url']}")
                if data.get("status_code"):
                    parts.append(f"Status: {data['status_code']}")
                text = " | ".join(parts)

            if not text:
                self.send_response(400)
                self._send_cors()
                self.send_header("Content-Type", "application/json")
                self.end_headers()
                self.wfile.write(json.dumps({"error": "Missing 'text' or telemetry fields"}).encode("utf-8"))
                return

            text = redact(text)
            key = hashlib.sha256(text.encode()).hexdigest()
            cached = _cache.get(key)
            if cached is not None:
                _stats["hits"] += 1
                out = {**cached, "cached": True}
                self.send_response(200)
                self._send_cors()
                self.send_header("Content-Type", "application/json")
                self.end_headers()
                self.wfile.write(json.dumps(out).encode("utf-8"))
                return
            _stats["misses"] += 1

            t_start = time.perf_counter()
            agent.reset()
            res = agent.run(text, strict=False)
            latency_ms = (time.perf_counter() - t_start) * 1000

            results = res.get("results", [])
            first_result = results[0] if results and isinstance(results[0], dict) else {}
            action = first_result.get("action", "needs_review")

            calls = res.get("function_calls", [])
            tool_name = calls[0].get("name") if calls else "triage_tool"

            out = {
                "success": True,
                "input_text": text,
                "action": action,
                "tool": tool_name,
                "arguments": first_result,
                "confidence": float(res.get("confidence", 0.0)),
                "reasoning": res.get("reasoning", ""),
                "latency_ms": round(latency_ms, 2),
                "model": "Cactus-Compute/needle3"
            }

            if len(_cache) >= CACHE_MAX:
                _cache.pop(next(iter(_cache)))
            _cache[key] = out

            self.send_response(200)
            self._send_cors()
            self.send_header("Content-Type", "application/json")
            self.end_headers()
            self.wfile.write(json.dumps(out).encode("utf-8"))
        else:
            self.send_response(404)
            self._send_cors()
            self.end_headers()

    def log_message(self, format, *args):
        # Concise logging
        print(f"[needle-sidecar] {args[0]} {args[1]} -> {args[2]}")


def run():
    server = ThreadingHTTPServer((HOST, PORT), NeedleHandler)
    print(f"[needle-sidecar] Listening on http://{HOST}:{PORT}")
    try:
        server.serve_forever()
    except KeyboardInterrupt:
        print("\nShutting down needle sidecar...")
        server.server_close()


if __name__ == "__main__":
    run()