atlas-chat 0.1.0__py3-none-any.whl

atlas/core/otel_config.py

@@ -0,0 +1,242 @@
+"""Unified logging & OpenTelemetry setup.
+
+Provides:
+- Structured JSON logging with optional trace/span identifiers
+- Environment or config-derived log level
+- Standard file output (project_root/logs/app.jsonl) with APP_LOG_DIR override
+- FastAPI & HTTPX instrumentation hooks
+"""
+
+from __future__ import annotations
+
+import json
+import logging
+import os
+from datetime import datetime, timezone
+from pathlib import Path
+from typing import Any, Dict, Optional
+
+from opentelemetry import trace
+from opentelemetry.instrumentation.fastapi import FastAPIInstrumentor
+from opentelemetry.instrumentation.httpx import HTTPXClientInstrumentor
+from opentelemetry.instrumentation.logging import LoggingInstrumentor
+from opentelemetry.sdk.resources import SERVICE_NAME, SERVICE_VERSION, Resource
+from opentelemetry.sdk.trace import TracerProvider
+
+
+class JSONFormatter(logging.Formatter):
+    """Format log records as JSON lines."""
+
+    def format(self, record: logging.LogRecord) -> str:  # noqa: D401
+        span = trace.get_current_span()
+        trace_id = span_id = None
+        if span and span.is_recording():
+            sc = span.get_span_context()
+            if sc.is_valid:
+                trace_id = f"{sc.trace_id:032x}"
+                span_id = f"{sc.span_id:016x}"
+
+        entry: Dict[str, Any] = {
+            "timestamp": datetime.fromtimestamp(record.created, tz=timezone.utc).isoformat(),
+            "level": record.levelname,
+            "logger": record.name,
+            "message": record.getMessage(),
+            "module": record.module,
+            "function": record.funcName,
+            "line": record.lineno,
+            "process_id": os.getpid(),
+            "thread_id": record.thread,
+            "thread_name": record.threadName,
+        }
+        if trace_id:
+            entry["trace_id"] = trace_id
+        if span_id:
+            entry["span_id"] = span_id
+        if record.exc_info:
+            entry["exception"] = self.formatException(record.exc_info)
+
+        excluded = {
+            "name","msg","args","levelname","levelno","pathname","filename","module","lineno",
+            "funcName","created","msecs","relativeCreated","thread","threadName","processName","process",
+            "exc_info","exc_text","stack_info","getMessage"
+        }
+        for k, v in record.__dict__.items():
+            if k not in excluded:
+                entry[f"extra_{k}"] = v
+        return json.dumps(entry, default=str)
+
+
+class OpenTelemetryConfig:
+    """Configure OpenTelemetry + structured logging."""
+
+    def __init__(self, service_name: str = "atlas-ui-3-backend", service_version: str = "1.0.0") -> None:
+        self.service_name = service_name
+        self.service_version = service_version
+        self.is_development = self._is_development()
+        self.log_level = self._get_log_level()
+        # Resolve logs directory robustly: use config manager
+        self.logs_dir = self._get_logs_dir()
+        self.log_file = self.logs_dir / "app.jsonl"
+        self.logs_dir.mkdir(parents=True, exist_ok=True)
+        self._setup_telemetry()
+        self._setup_logging()
+
+    # ------------------------------------------------------------------
+    # Internals
+    # ------------------------------------------------------------------
+    def _get_logs_dir(self) -> Path:
+        """Get logs directory from config manager or default to project_root/logs."""
+        try:
+            from atlas.modules.config import config_manager
+            if config_manager.app_settings.app_log_dir:
+                return Path(config_manager.app_settings.app_log_dir)
+        except Exception:
+            # Config manager may not be initialized during early startup or tests.
+            # Fall back to default logs directory without logging (avoid circular deps).
+            pass
+        # Fallback: project_root/logs
+        project_root = Path(__file__).resolve().parents[2]
+        return project_root / "logs"
+
+    def _is_development(self) -> bool:
+        try:
+            from atlas.modules.config import config_manager
+            settings = config_manager.app_settings
+            return (
+                settings.debug_mode
+                or settings.environment.lower() in {"dev", "development"}
+            )
+        except Exception:
+            # Fallback to environment variables if config not available
+            return (
+                os.getenv("DEBUG_MODE", "false").lower() == "true"
+                or os.getenv("ENVIRONMENT", "production").lower() in {"dev", "development"}
+            )
+
+    def _get_log_level(self) -> int:
+        try:
+            from atlas.modules.config import config_manager
+            level_name = config_manager.app_settings.log_level.upper()
+        except Exception:
+            # Fallback to environment variable if config not available
+            level_name = os.getenv("LOG_LEVEL", "INFO").upper()
+        level = getattr(logging, level_name, None)
+        return level if isinstance(level, int) else logging.INFO
+
+    def _setup_telemetry(self) -> None:
+        resource = Resource.create(
+            {
+                SERVICE_NAME: self.service_name,
+                SERVICE_VERSION: self.service_version,
+                "environment": "development" if self.is_development else "production",
+            }
+        )
+        trace.set_tracer_provider(TracerProvider(resource=resource))
+
+    def _setup_logging(self) -> None:
+        root = logging.getLogger()
+        for h in root.handlers[:]:
+            root.removeHandler(h)
+
+        json_formatter = JSONFormatter()
+        file_handler = logging.FileHandler(self.log_file, encoding="utf-8")
+        file_handler.setFormatter(json_formatter)
+        file_handler.setLevel(self.log_level)
+        root.addHandler(file_handler)
+        root.setLevel(self.log_level)
+
+        # Reduce noise from third-party libraries at INFO.
+        # We still want their warnings/errors, and their debug output remains available
+        # when LOG_LEVEL=DEBUG.
+        if self.log_level > logging.DEBUG:
+            for noisy in (
+                "httpx",
+                "httpcore",
+                "LiteLLM",
+                "litellm",
+            ):
+                logging.getLogger(noisy).setLevel(logging.WARNING)
+
+        if self.is_development:
+            console = logging.StreamHandler()
+            console.setFormatter(logging.Formatter("%(asctime)s - %(name)s - %(levelname)s - %(message)s"))
+            console.setLevel(logging.WARNING)
+            root.addHandler(console)
+            for noisy in (
+                "httpx",
+                "urllib3.connectionpool",
+                "auth_utils",
+                "message_processor",
+                "session",
+                "callbacks",
+                "utils",
+                "banner_client",
+                "middleware",
+                "mcp_client",
+            ):
+                logging.getLogger(noisy).setLevel(logging.DEBUG)
+
+        LoggingInstrumentor().instrument(set_logging_format=False)
+
+    # ------------------------------------------------------------------
+    # Public helpers
+    # ------------------------------------------------------------------
+    def instrument_fastapi(self, app) -> None:  # noqa: ANN001
+        FastAPIInstrumentor.instrument_app(app)
+
+    def instrument_httpx(self) -> None:
+        HTTPXClientInstrumentor().instrument()
+
+    def get_log_file_path(self) -> Path:
+        return self.log_file
+
+    def read_logs(self, lines: int = 100) -> list[Dict[str, Any]]:
+        if not self.log_file.exists():
+            return []
+        out: list[Dict[str, Any]] = []
+        try:
+            with self.log_file.open("r", encoding="utf-8") as f:
+                data = f.readlines()[-lines:]
+            for ln in data:
+                ln = ln.strip()
+                if not ln:
+                    continue
+                try:
+                    out.append(json.loads(ln))
+                except json.JSONDecodeError:
+                    continue
+        except Exception as e:  # noqa: BLE001
+            logging.getLogger(__name__).error(f"Error reading logs: {e}")
+        return out
+
+    def get_log_stats(self) -> Dict[str, Any]:
+        if not self.log_file.exists():
+            return {"file_exists": False, "file_size": 0, "line_count": 0, "last_modified": None}
+        try:
+            stat = self.log_file.stat()
+            with self.log_file.open("r", encoding="utf-8") as f:
+                line_count = sum(1 for _ in f)
+            return {
+                "file_exists": True,
+                "file_size": stat.st_size,
+                "line_count": line_count,
+                "last_modified": datetime.fromtimestamp(stat.st_mtime, tz=timezone.utc).isoformat(),
+                "file_path": str(self.log_file),
+            }
+        except Exception as e:  # noqa: BLE001
+            logging.getLogger(__name__).error(f"Error getting log stats: {e}")
+            return {"file_exists": True, "error": str(e)}
+
+
+# Global instance
+otel_config: Optional[OpenTelemetryConfig] = None
+
+
+def setup_opentelemetry(service_name: str = "atlas-ui-3-backend", service_version: str = "1.0.0") -> OpenTelemetryConfig:
+    global otel_config
+    otel_config = OpenTelemetryConfig(service_name, service_version)
+    return otel_config
+
+
+def get_otel_config() -> Optional[OpenTelemetryConfig]:
+    return otel_config

atlas/core/prompt_risk.py

@@ -0,0 +1,200 @@
+"""
+Prompt injection risk heuristics and structured logging.
+
+Scope: lightweight, configurable thresholds; used for user input and RAG results.
+"""
+
+from __future__ import annotations
+
+import base64
+import json
+import logging
+import math
+import re
+from collections import Counter
+from datetime import datetime, timezone
+from pathlib import Path
+from typing import Dict, List, Optional
+
+logger = logging.getLogger(__name__)
+
+
+def _get_thresholds() -> Dict[str, int]:
+    """Get prompt injection risk thresholds from config manager."""
+    try:
+        from atlas.modules.config import config_manager
+        settings = config_manager.app_settings
+        return {
+            "low": settings.pi_threshold_low,
+            "medium": settings.pi_threshold_medium,
+            "high": settings.pi_threshold_high,
+        }
+    except Exception:
+        # Fallback to defaults if config not available
+        return {
+            "low": 30,
+            "medium": 50,
+            "high": 80,
+        }
+
+
+def calculate_prompt_injection_risk(message: str, *, mode: str = "general") -> Dict[str, object]:
+    """
+    Calculate a heuristic risk score for prompt injection attempts.
+
+    Returns: { 'score': int, 'risk_level': str, 'triggers': list[str] }
+    """
+    score = 0
+    triggers: List[str] = []
+
+    msg_lower = (message or "").lower()
+
+    # 1) Suspicious patterns
+    patterns = {
+        "override_instructions": (r"ignore\s+(previous|all|everything|above|prior)", 40),
+        "disregard": (r"disregard\s+(previous|all|everything|above|prior)", 40),
+        "new_instructions": (r"new\s+instructions?\s*:\s*", 35),
+        "system_role": (r"\b(system|assistant|user)\s*:\s*", 30),
+        "act_as": (r"act\s+as\s+(if\s+)?you\s+(are|were)", 25),
+        "pretend": (r"pretend\s+(to\s+be|you\s+are)", 25),
+        "role_change": (r"your?\s+(new\s+)?role\s+(is|now)", 30),
+        "forget": (r"forget\s+(everything|all|previous)", 35),
+        "override": (r"override\s+(previous|default|system)", 35),
+        "jailbreak": (r"(jailbreak|developer\s+mode|god\s+mode)", 45),
+    }
+    for name, (pat, pts) in patterns.items():
+        if re.search(pat, msg_lower):
+            score += pts
+            triggers.append(name)
+
+    # 2) Encodings/obfuscation
+    if _detect_encoding(message):
+        score += 30
+        triggers.append("encoding_detected")
+
+    # 3) Statistical anomalies
+    # Delimiter density (triple quotes, fences, etc.)
+    delimiters = len(re.findall(r"[#*\-_=]{3,}|[\"\"\"''']{3,}", message or ""))
+    if delimiters >= 3:
+        score += 25
+        triggers.append("excessive_delimiters")
+    elif delimiters >= 1:
+        score += 10
+
+    # High entropy (possible encoded blob)
+    if len(message or "") > 10:
+        ent = _calculate_entropy(message)
+        if ent > 4.5:
+            score += 20
+            triggers.append("high_entropy")
+
+    # Excessive caps
+    if len(message or "") > 20:
+        caps_ratio = sum(1 for c in (message or "") if c.isupper()) / max(1, len(message or ""))
+        if caps_ratio > 0.3:
+            score += 15
+            triggers.append("excessive_caps")
+
+    # 4) Context-breaking attempts
+    if (message or "").count("\n") > 5 or re.search(r"\s{10,}", message or ""):
+        score += 15
+        triggers.append("formatting_abuse")
+
+    if re.search(r"(human|user|assistant):\s*\n", msg_lower):
+        score += 25
+        triggers.append("fake_conversation")
+
+    if (len(message or "") > 50) and re.search(r"<[^>]+>.*</[^>]+>|[{}\[\]]", message or ""):
+        score += 20
+        triggers.append("structured_injection")
+
+    # 5) Length penalty
+    if len(message or "") > 1000:
+        score += 15
+        triggers.append("excessive_length")
+
+    # Context-aware normalization (reduce false positives)
+    if mode in ("code", "logs"):
+        # Code/logs: braces, fences common; soften penalties
+        if "excessive_delimiters" in triggers:
+            score -= 10
+        if "structured_injection" in triggers:
+            score -= 10
+        score = max(0, score)
+
+    # Risk buckets - get thresholds from config
+    thresholds = _get_thresholds()
+    if score >= thresholds["high"]:
+        level = "high"
+    elif score >= thresholds["medium"]:
+        level = "medium"
+    elif score >= thresholds["low"]:
+        level = "low"
+    else:
+        level = "minimal"
+
+    return {"score": int(score), "risk_level": level, "triggers": triggers}
+
+
+def _detect_encoding(text: str) -> bool:
+    clean = re.sub(r"\s", "", text or "")
+    # Base64-like
+    if len(clean) > 20 and len(clean) % 4 == 0 and re.match(r"^[A-Za-z0-9+/=]+$", clean or ""):
+        try:
+            base64.b64decode(clean, validate=True)
+            return True
+        except Exception:
+            pass
+    # Hex
+    if re.match(r"^(0x)?[0-9a-fA-F]+$", clean or "") and len(clean) > 20:
+        return True
+    # Escape sequences
+    if re.search(r"\\u[0-9a-fA-F]{4}|\\x[0-9a-fA-F]{2}", text or ""):
+        return True
+    # Zero-width/unusual unicode
+    if re.search(r"[\u200B-\u200D\uFEFF\u2060]", text or ""):
+        return True
+    return False
+
+
+def _calculate_entropy(text: str) -> float:
+    if not text:
+        return 0.0
+    counts = Counter(text)
+    n = len(text)
+    ent = 0.0
+    for c in counts.values():
+        p = c / n
+        ent -= p * math.log2(max(p, 1e-12))
+    return ent
+
+
+def log_high_risk_event(*, source: str, user: Optional[str], content: str, score: int, risk_level: str, triggers: List[str], extra: Optional[Dict[str, object]] = None) -> None:
+    """Append a JSONL record for medium/high events to logs/security_high_risk.jsonl."""
+    try:
+        # Only log medium/high
+        if risk_level not in ("medium", "high"):
+            return
+        base_dir = Path(__file__).resolve().parents[2]
+        log_path = base_dir / "logs" / "security_high_risk.jsonl"
+        log_path.parent.mkdir(parents=True, exist_ok=True)
+        record = {
+            "ts": datetime.now(timezone.utc).isoformat() + "Z",
+            "type": "prompt_risk",
+            "source": source,
+            "user": user,
+            "score": score,
+            "risk_level": risk_level,
+            "triggers": triggers,
+        }
+        if extra:
+            record.update(extra)
+        # include a small snippet only
+        snippet = (content or "")
+        if len(snippet) > 240:
+            snippet = snippet[:240] + "…"
+        record["snippet"] = snippet
+        with open(log_path, "a", encoding="utf-8") as f:
+            f.write(json.dumps(record, ensure_ascii=False) + "\n")
+    except Exception as e:
+        logger.debug("Failed to write high risk log: %s", e)

atlas/core/rate_limit_middleware.py

@@ -0,0 +1,64 @@
+"""Simple in-memory rate limit middleware.
+
+Fixed-window counter per client IP (and optionally per-path) to throttle requests.
+This is a lightweight safeguard suitable for single-process deployments and tests.
+
+Configuration is sourced from ConfigManager (AppSettings) with optional env overrides:
+    - app_settings.rate_limit_rpm            (env: RATE_LIMIT_RPM, default: 600)
+    - app_settings.rate_limit_window_seconds (env: RATE_LIMIT_WINDOW_SECONDS, default: 60)
+    - app_settings.rate_limit_per_path       (env: RATE_LIMIT_PER_PATH, default: false)
+"""
+
+import time
+import typing as t
+
+from fastapi import Request
+from starlette.middleware.base import BaseHTTPMiddleware
+from starlette.responses import JSONResponse, Response
+
+from atlas.modules.config import config_manager
+
+
+class RateLimitMiddleware(BaseHTTPMiddleware):
+    def __init__(self, app) -> None:
+        super().__init__(app)
+        settings = config_manager.app_settings
+        # Pull from centralized config with sane defaults
+        self.window_seconds = int(getattr(settings, "rate_limit_window_seconds", 60))
+        self.max_requests = int(getattr(settings, "rate_limit_rpm", 600))
+        self.per_path = bool(getattr(settings, "rate_limit_per_path", False))
+        # state: key -> (window_start_epoch, count)
+        self._buckets: dict[str, t.Tuple[int, int]] = {}
+
+    def _key_for(self, request: Request) -> str:
+        client_ip = getattr(request.client, "host", "unknown") if request.client else "unknown"
+        if self.per_path:
+            return f"{client_ip}:{request.url.path}"
+        return client_ip
+
+    async def dispatch(self, request: Request, call_next) -> Response:
+        now = int(time.time())
+        key = self._key_for(request)
+        win = self.window_seconds
+        start, count = self._buckets.get(key, (now, 0))
+
+        # Move window if expired
+        if now - start >= win:
+            start, count = now, 0
+
+        count += 1
+        self._buckets[key] = (start, count)
+
+        if count > self.max_requests:
+            retry_after = max(1, win - (now - start))
+            return JSONResponse(
+                status_code=429,
+                content={
+                    "detail": "Rate limit exceeded. Please try again later.",
+                    "limit": self.max_requests,
+                    "window_seconds": self.window_seconds,
+                },
+                headers={"Retry-After": str(retry_after)},
+            )
+
+        return await call_next(request)

atlas/core/security_headers_middleware.py

@@ -0,0 +1,51 @@
+"""Security headers middleware with ConfigManager-based toggles.
+
+Sets common security headers:
+ - Content-Security-Policy (CSP)
+ - X-Frame-Options (XFO)
+ - X-Content-Type-Options: nosniff
+ - Referrer-Policy
+
+Each header is individually togglable via AppSettings. HSTS is intentionally omitted.
+"""
+
+from fastapi import Request
+from starlette.middleware.base import BaseHTTPMiddleware
+from starlette.responses import Response
+from starlette.types import ASGIApp
+
+from atlas.modules.config import config_manager
+
+
+class SecurityHeadersMiddleware(BaseHTTPMiddleware):
+    def __init__(self, app: ASGIApp) -> None:
+        super().__init__(app)
+        self.settings = config_manager.app_settings
+
+    async def dispatch(self, request: Request, call_next):
+        response: Response = await call_next(request)
+
+        # X-Content-Type-Options
+        if getattr(self.settings, "security_nosniff_enabled", True):
+            if "X-Content-Type-Options" not in response.headers:
+                response.headers["X-Content-Type-Options"] = "nosniff"
+
+        # X-Frame-Options
+        if getattr(self.settings, "security_xfo_enabled", True):
+            xfo_value = getattr(self.settings, "security_xfo_value", "SAMEORIGIN")
+            if "X-Frame-Options" not in response.headers:
+                response.headers["X-Frame-Options"] = xfo_value
+
+        # Referrer-Policy
+        if getattr(self.settings, "security_referrer_policy_enabled", True):
+            ref_value = getattr(self.settings, "security_referrer_policy_value", "no-referrer")
+            if "Referrer-Policy" not in response.headers:
+                response.headers["Referrer-Policy"] = ref_value
+
+        # Content-Security-Policy
+        if getattr(self.settings, "security_csp_enabled", True):
+            csp_value = getattr(self.settings, "security_csp_value", None)
+            if csp_value and "Content-Security-Policy" not in response.headers:
+                response.headers["Content-Security-Policy"] = csp_value
+
+        return response

atlas/domain/__init__.py

@@ -0,0 +1,37 @@
+"""Domain layer - pure business models and logic."""
+
+from .errors import (
+    AuthenticationError,
+    AuthorizationError,
+    ConfigurationError,
+    DomainError,
+    LLMError,
+    MessageError,
+    SessionError,
+    ToolError,
+    ValidationError,
+)
+from .messages.models import ConversationHistory, Message, MessageRole, MessageType, ToolCall, ToolResult
+from .sessions.models import Session
+
+__all__ = [
+    # Errors
+    "DomainError",
+    "ValidationError",
+    "SessionError",
+    "MessageError",
+    "AuthenticationError",
+    "AuthorizationError",
+    "ConfigurationError",
+    "LLMError",
+    "ToolError",
+    # Messages
+    "Message",
+    "MessageRole",
+    "MessageType",
+    "ToolCall",
+    "ToolResult",
+    "ConversationHistory",
+    # Sessions
+    "Session",
+]

atlas/domain/chat/__init__.py

@@ -0,0 +1 @@
+"""Domain models for chat operations."""

atlas/domain/chat/dtos.py

@@ -0,0 +1,85 @@
+"""Data Transfer Objects for chat operations."""
+
+from dataclasses import dataclass, field
+from typing import Any, Dict, List, Optional
+from uuid import UUID
+
+
+@dataclass
+class ChatRequest:
+    """
+    Request DTO for chat operations.
+
+    Contains all parameters needed for different chat modes (plain, tools, RAG, agent).
+    """
+    session_id: UUID
+    content: str
+    model: str
+    user_email: Optional[str] = None
+    selected_tools: Optional[List[str]] = None
+    selected_prompts: Optional[List[str]] = None
+    selected_data_sources: Optional[List[str]] = None
+    only_rag: bool = False
+    tool_choice_required: bool = False
+    agent_mode: bool = False
+    temperature: float = 0.7
+    agent_max_steps: int = 30
+    agent_loop_strategy: Optional[str] = None
+    files: Optional[Dict[str, Any]] = None
+    extra: Dict[str, Any] = field(default_factory=dict)
+
+
+@dataclass
+class ChatResponse:
+    """
+    Response DTO for chat operations.
+
+    Contains the result of a chat interaction.
+    """
+    type: str
+    message: str
+    metadata: Dict[str, Any] = field(default_factory=dict)
+
+    def to_dict(self) -> Dict[str, Any]:
+        """Convert to dictionary format for API response."""
+        return {
+            "type": self.type,
+            "message": self.message,
+            **self.metadata
+        }
+
+
+@dataclass
+class LLMMessage:
+    """
+    Type-safe message format for LLM interactions.
+
+    Normalizes message structure across different chat modes.
+    """
+    role: str  # "user", "assistant", "system", "tool"
+    content: str
+    name: Optional[str] = None
+    tool_calls: Optional[List[Dict[str, Any]]] = None
+    tool_call_id: Optional[str] = None
+
+    def to_dict(self) -> Dict[str, Any]:
+        """Convert to dictionary format for LLM API."""
+        result = {"role": self.role, "content": self.content}
+        if self.name:
+            result["name"] = self.name
+        if self.tool_calls:
+            result["tool_calls"] = self.tool_calls
+        if self.tool_call_id:
+            result["tool_call_id"] = self.tool_call_id
+        return result
+
+    @classmethod
+    def from_dict(cls, data: Dict[str, Any]) -> "LLMMessage":
+        """Create from dictionary format."""
+        return cls(
+            role=data["role"],
+            content=data.get("content", ""),
+            name=data.get("name"),
+            tool_calls=data.get("tool_calls"),
+            tool_call_id=data.get("tool_call_id"),
+        )