diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000000000000000000000000000000000000..476a2e53de797a3668bec5f37e9d70caab50ad9d --- /dev/null +++ b/.gitignore @@ -0,0 +1,7 @@ +__pycache__/ +*.py[cod] +venv/ +.env +*.log +.pytest_cache/ +test.db diff --git a/Dockerfile b/Dockerfile new file mode 100644 index 0000000000000000000000000000000000000000..02c8bca0f7a51bee38ba16708c20c09c93171a76 --- /dev/null +++ b/Dockerfile @@ -0,0 +1,24 @@ +FROM python:3.11-slim + +WORKDIR /app + +RUN apt-get update && apt-get install -y \ + libpq-dev gcc libmagic1 curl \ + && rm -rf /var/lib/apt/lists/* + +RUN pip install --no-cache-dir torch --index-url https://download.pytorch.org/whl/cpu + +COPY requirements.txt . +RUN pip install --no-cache-dir -r requirements.txt + +RUN python -c "from sentence_transformers import CrossEncoder; CrossEncoder('cross-encoder/ms-marco-MiniLM-L-6-v2')" + +ENV TOKENIZERS_PARALLELISM=false +ENV TRANSFORMERS_CACHE=/app/.cache/huggingface +ENV SENTENCE_TRANSFORMERS_HOME=/app/.cache/sentence-transformers + +COPY . . + +EXPOSE 8000 + +CMD ["uvicorn", "app.main:app", "--host", "0.0.0.0", "--port", "8000"] diff --git a/README.md b/README.md index a2a005ab726e731fe6a4e86e71d33077bbc4d856..5f68146e716c356da9c49ca06125427360e7a2eb 100644 --- a/README.md +++ b/README.md @@ -1,11 +1,19 @@ --- -title: Miningniti Api -emoji: 🐨 -colorFrom: yellow -colorTo: pink +title: MiningNiti API +emoji: ⛏️ +colorFrom: purple +colorTo: blue sdk: docker -pinned: false -license: apache-2.0 +app_port: 8000 +pinned: true --- -Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference +# MiningNiti API + +AI-powered document intelligence engine for the coal mining industry. + +## Features +- Multi-agent AI pipeline (6 agents, 4 providers) +- Production RAG with hybrid search + cross-encoder reranking +- Compliance auto-auditing +- Real-time streaming chat diff --git a/alembic.ini b/alembic.ini new file mode 100644 index 0000000000000000000000000000000000000000..2304f6f1b604985a98f33903183cc53274bdae40 --- /dev/null +++ b/alembic.ini @@ -0,0 +1,112 @@ +# A generic, single database configuration. + +[alembic] +# path to migration scripts +script_location = alembic + +# template used to generate migration file names; The default value is %%(rev)s_%%(slug)s +file_template = %%(rev)s_%%(slug)s + +# sys.path path, will be prepended to sys.path if present. +# defaults to the current working directory. +prepend_sys_path = . + +# timezone to use when rendering the date within the migration file +# as well as the filename. +# If specified, requires the python>=3.9 or backports.zoneinfo library. +# Any required deps can installed via pip[tz] +# timezone = UTC + +# max length of characters to apply to the "slug" field +# truncate_slug_length = 40 + +# set to 'true' to run the environment during +# the 'revision' command, regardless of autogenerate +# revision_environment = false + +# set to 'true' to allow .pyc and .pyo files without +# a source .py file to be detected as revisions in the +# versions/ directory +# sourceless = false + +# version location specification; This defaults +# to alembic/versions. When using multiple version +# directories, initial revisions must be specified with --version-path. +# The path separator used here should be the separator specified by "version_path_separator" below. +# version_locations = %(here)s/bar:%(here)s/bat:alembic/versions + +# version path separator; As mentioned above, this is the character used to split +# version_locations. The default within new alembic.ini files is "os", which uses os.pathsep. +# If this key is omitted entirely, it falls back to the legacy behavior of splitting on spaces and/or commas. +# Valid values for version_path_separator are: +# +# version_path_separator = : +# version_path_separator = ; +# version_path_separator = space +version_path_separator = os # Use os.pathsep. Default configuration used for new projects. + +# set to 'true' to search source files recursively +# in each "version_locations" directory +# New in Alembic version 1.10 +# recursive_version_locations = false + +# the output encoding used when revision files +# are written from script.py.mako +# output_encoding = utf-8 + +# DATABASE_URL is loaded from environment variable in env.py +# Do NOT hardcode the URL here — use .env file +sqlalchemy.url = driver://user:pass@localhost/dbname + + +[post_write_hooks] +# post_write_hooks defines scripts or Python functions that are run +# on newly generated revision scripts. See the documentation for further +# detail and examples + +# format using "black" - use the console_scripts runner, against the "black" entrypoint +# hooks = black +# black.type = console_scripts +# black.entrypoint = black +# black.options = -l 79 REVISION_SCRIPT_FILENAME + +# lint with attempts to fix using "ruff" - use the exec runner, against a script +# hooks = ruff +# ruff.type = exec +# ruff.executable = %(here)s/.venv/bin/ruff +# ruff.options = --fix REVISION_SCRIPT_FILENAME + +# Logging configuration +[loggers] +keys = root,sqlalchemy,alembic + +[handlers] +keys = console + +[formatters] +keys = generic + +[logger_root] +level = WARNING +handlers = console +qualname = + +[logger_sqlalchemy] +level = WARNING +handlers = +qualname = sqlalchemy.engine + +[logger_alembic] +level = INFO +handlers = +qualname = alembic + +[handler_console] +class = StreamHandler +args = (sys.stderr,) +level = NOTSET +formatter = generic + +[formatter_generic] +format = %(levelname)-5.5s [%(name)s] %(message)s +datefmt = %H:%M:%S diff --git a/alembic/README b/alembic/README new file mode 100644 index 0000000000000000000000000000000000000000..98e4f9c44effe479ed38c66ba922e7bcc672916f --- /dev/null +++ b/alembic/README @@ -0,0 +1 @@ +Generic single-database configuration. \ No newline at end of file diff --git a/alembic/env.py b/alembic/env.py new file mode 100644 index 0000000000000000000000000000000000000000..7fb7ec8065f533ddc3babc279f2ec71a3f676987 --- /dev/null +++ b/alembic/env.py @@ -0,0 +1,84 @@ +""" +Alembic Environment Configuration +Loads DATABASE_URL from .env and connects all SQLAlchemy models. +""" + +import os +import sys +from logging.config import fileConfig +from pathlib import Path + +from sqlalchemy import engine_from_config, pool, text +from alembic import context + +# ── Ensure 'backend/' is on sys.path so `app.*` imports work ────────────────── +BACKEND_DIR = Path(__file__).resolve().parent.parent +sys.path.insert(0, str(BACKEND_DIR)) + +# ── Load settings (reads .env automatically via pydantic-settings) ───────────── +from app.config import settings # noqa: E402 + +# ── Import ALL models so Alembic can detect their tables ────────────────────── +from app.models.base import Base # noqa: E402 +from app.models import user, document, chat, audit, prompt # noqa: E402, F401 + +# ── Alembic Config object ───────────────────────────────────────────────────── +config = context.config + +# Override sqlalchemy.url with the value from our settings/.env +config.set_main_option("sqlalchemy.url", settings.DATABASE_URL) + +# Interpret the config file for Python logging (if present) +if config.config_file_name is not None: + fileConfig(config.config_file_name) + +# Target metadata for autogenerate support +target_metadata = Base.metadata + + +def run_migrations_offline() -> None: + """ + Run migrations in 'offline' mode. + Generates SQL scripts without a live DB connection. + """ + url = config.get_main_option("sqlalchemy.url") + context.configure( + url=url, + target_metadata=target_metadata, + literal_binds=True, + dialect_opts={"paramstyle": "named"}, + compare_type=True, + ) + + with context.begin_transaction(): + context.run_migrations() + + +def run_migrations_online() -> None: + """ + Run migrations in 'online' mode. + Connects to the real database and applies migrations. + """ + connectable = engine_from_config( + config.get_section(config.config_ini_section, {}), + prefix="sqlalchemy.", + poolclass=pool.NullPool, + ) + + with connectable.connect() as connection: + context.configure( + connection=connection, + target_metadata=target_metadata, + compare_type=True, + # Include schemas for pgvector extension objects + include_schemas=False, + ) + + with context.begin_transaction(): + context.run_migrations() + + +if context.is_offline_mode(): + run_migrations_offline() +else: + run_migrations_online() diff --git a/alembic/script.py.mako b/alembic/script.py.mako new file mode 100644 index 0000000000000000000000000000000000000000..fbc4b07dcef98b20c6f96b642097f35e8433258e --- /dev/null +++ b/alembic/script.py.mako @@ -0,0 +1,26 @@ +"""${message} + +Revision ID: ${up_revision} +Revises: ${down_revision | comma,n} +Create Date: ${create_date} + +""" +from typing import Sequence, Union + +from alembic import op +import sqlalchemy as sa +${imports if imports else ""} + +# revision identifiers, used by Alembic. +revision: str = ${repr(up_revision)} +down_revision: Union[str, None] = ${repr(down_revision)} +branch_labels: Union[str, Sequence[str], None] = ${repr(branch_labels)} +depends_on: Union[str, Sequence[str], None] = ${repr(depends_on)} + + +def upgrade() -> None: + ${upgrades if upgrades else "pass"} + + +def downgrade() -> None: + ${downgrades if downgrades else "pass"} diff --git a/alembic/versions/001_enable_pgvector.py b/alembic/versions/001_enable_pgvector.py new file mode 100644 index 0000000000000000000000000000000000000000..7318b3a2f7370c24b03bd90be591302abd1e2d1a --- /dev/null +++ b/alembic/versions/001_enable_pgvector.py @@ -0,0 +1,153 @@ +""" +Migration: Enable pgvector and migrate embeddings from JSONB to vector(768) + +Revision ID: 001 +Revises: (initial) +Create Date: 2026-06-07 + +Changes: + 1. Enable the pgvector extension + 2. Add page tracking columns to document_embeddings + (section_title, page_numbers) + 3. Migrate embedding column from JSONB to vector(768) native type + 4. Create HNSW index for fast ANN search (~5ms vs seconds) + 5. Fix datetime.utcnow() deprecation in base models (add timezone info) + 6. Add total_pages column to documents +""" + +from alembic import op +import sqlalchemy as sa +from sqlalchemy.dialects import postgresql + +# revision identifiers, used by Alembic +revision = "001" +down_revision = None +branch_labels = None +depends_on = None + + +def upgrade() -> None: + # ── Step 1: Enable pgvector extension ────────────────────────────────────── + op.execute("CREATE EXTENSION IF NOT EXISTS vector") + op.execute("CREATE EXTENSION IF NOT EXISTS pg_trgm") # For full-text search + + # ── Step 2: Add page tracking columns to document_embeddings ─────────────── + op.add_column( + "document_embeddings", + sa.Column("section_title", sa.String(500), nullable=True), + ) + op.add_column( + "document_embeddings", + sa.Column( + "page_numbers", + postgresql.JSONB(astext_type=sa.Text()), + nullable=True, + comment="List of page numbers this chunk spans, e.g. [12, 13]", + ), + ) + + # ── Step 3: Add total_pages column to documents ──────────────────────────── + # (page_count already exists — we just ensure it's correctly named) + # Add total_pages as an alias; keep page_count for backward compatibility + op.add_column( + "documents", + sa.Column("total_pages", sa.Integer(), nullable=True), + ) + + # ── Step 4: Migrate embedding column from JSONB → vector(768) ───────────── + # First add the new column + op.execute( + "ALTER TABLE document_embeddings ADD COLUMN embedding_vec vector(768)" + ) + + # Convert existing JSONB embeddings to vector type + # This handles both list-of-floats and null values safely + op.execute( + """ + UPDATE document_embeddings + SET embedding_vec = ( + SELECT array_agg(elem::float8)::vector(768) + FROM jsonb_array_elements_text(embedding) AS elem + ) + WHERE embedding IS NOT NULL + AND jsonb_typeof(embedding) = 'array' + AND jsonb_array_length(embedding) = 768 + """ + ) + + # Drop the old JSONB column and rename new one + op.execute("ALTER TABLE document_embeddings DROP COLUMN embedding") + op.execute( + "ALTER TABLE document_embeddings RENAME COLUMN embedding_vec TO embedding" + ) + + # Make embedding NOT NULL (existing rows already converted) + op.execute( + "ALTER TABLE document_embeddings ALTER COLUMN embedding SET NOT NULL" + ) + + # ── Step 5: Create HNSW index for approximate nearest neighbor search ────── + # HNSW gives sub-5ms search up to ~1M vectors + # m=16: max connections per node (higher = better recall, more memory) + # ef_construction=200: build-time search depth (higher = better quality index) + op.execute( + """ + CREATE INDEX idx_embeddings_hnsw + ON document_embeddings + USING hnsw (embedding vector_cosine_ops) + WITH (m = 16, ef_construction = 200) + """ + ) + + # ── Step 6: Add composite index on (document_id, chunk_index) ────────────── + op.create_index( + "idx_embeddings_doc_chunk", + "document_embeddings", + ["document_id", "chunk_index"], + unique=True, + ) + + # ── Step 7: Add trigram index on documents for full-text search ──────────── + op.execute( + """ + CREATE INDEX idx_documents_title_trgm + ON documents + USING gin (title gin_trgm_ops) + """ + ) + op.execute( + """ + CREATE INDEX idx_documents_filename_trgm + ON documents + USING gin (file_name gin_trgm_ops) + """ + ) + + +def downgrade() -> None: + # Remove indexes + op.execute("DROP INDEX IF EXISTS idx_embeddings_hnsw") + op.execute("DROP INDEX IF EXISTS idx_documents_title_trgm") + op.execute("DROP INDEX IF EXISTS idx_documents_filename_trgm") + op.drop_index("idx_embeddings_doc_chunk", table_name="document_embeddings") + + # Restore JSONB column + op.execute( + "ALTER TABLE document_embeddings ADD COLUMN embedding_jsonb jsonb" + ) + op.execute( + """ + UPDATE document_embeddings + SET embedding_jsonb = to_jsonb(embedding::float8[]) + WHERE embedding IS NOT NULL + """ + ) + op.execute("ALTER TABLE document_embeddings DROP COLUMN embedding") + op.execute( + "ALTER TABLE document_embeddings RENAME COLUMN embedding_jsonb TO embedding" + ) + + # Remove added columns + op.drop_column("document_embeddings", "section_title") + op.drop_column("document_embeddings", "page_numbers") + op.drop_column("documents", "total_pages") diff --git a/alembic/versions/002_hybrid_search_index.py b/alembic/versions/002_hybrid_search_index.py new file mode 100644 index 0000000000000000000000000000000000000000..080cbe3f916c174c2d5aac57bf4dcd98874e94c4 --- /dev/null +++ b/alembic/versions/002_hybrid_search_index.py @@ -0,0 +1,43 @@ +"""Add GIN trigram index on chunk_text for hybrid search + +Revision ID: 002 +Revises: 001_enable_pgvector +Create Date: 2025-01-01 +""" + +from alembic import op +import sqlalchemy as sa + +revision = "002" +down_revision = "001" +branch_labels = None +depends_on = None + + +def upgrade() -> None: + # Ensure pg_trgm extension exists (for trigram similarity) + op.execute("CREATE EXTENSION IF NOT EXISTS pg_trgm") + + # GIN trigram index on chunk_text for fast BM25-style keyword search + # Enables: WHERE chunk_text % :query (trigram similarity match) + op.execute( + """ + CREATE INDEX IF NOT EXISTS idx_embeddings_chunk_text_trgm + ON document_embeddings + USING gin (chunk_text gin_trgm_ops) + """ + ) + + # GIN trigram index on document content for full-document keyword search + op.execute( + """ + CREATE INDEX IF NOT EXISTS idx_documents_content_trgm + ON documents + USING gin (content gin_trgm_ops) + """ + ) + + +def downgrade() -> None: + op.execute("DROP INDEX IF EXISTS idx_embeddings_chunk_text_trgm") + op.execute("DROP INDEX IF EXISTS idx_documents_content_trgm") diff --git a/app/README.md b/app/README.md new file mode 100644 index 0000000000000000000000000000000000000000..2885772e7590e13c100abd654cc7d5a63f36f8a9 --- /dev/null +++ b/app/README.md @@ -0,0 +1,30 @@ +# MiningNiti Enterprise Backend + +Production-ready AI Document Intelligence Engine for the Mining Industry. + +## Structure + +``` +app/ +├── api/ # REST API endpoints +├── agents/ # AI agents (LangGraph) +├── core/ # Security, config, exceptions +├── db/ # Database session & migrations +├── models/ # SQLAlchemy ORM models +├── schemas/ # Pydantic request/response schemas +├── services/ # Business logic layer +└── workers/ # Celery background tasks +``` + +## Quick Start + +```bash +# Install dependencies +pip install -r requirements.txt + +# Run development server +uvicorn app.main:app --reload --port 8000 + +# Run Celery worker +celery -A app.workers.celery_app worker -l info +``` diff --git a/app/__init__.py b/app/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..b7e1c00601eca77a0838b244c776d0ecf783629b --- /dev/null +++ b/app/__init__.py @@ -0,0 +1,7 @@ +""" +MiningNiti Enterprise Backend +AI-Powered Document Intelligence for the Mining Industry +""" + +__version__ = "2.0.0" +__author__ = "MiningNiti Team" diff --git a/app/agents/__init__.py b/app/agents/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..516d18440478617c88937176cdf66ff598263d3d --- /dev/null +++ b/app/agents/__init__.py @@ -0,0 +1,20 @@ +""" +AI Agents Module +Specialized agents for mining document intelligence +""" + +from app.agents.base import BaseAgent +from app.agents.classifier import ClassifierAgent +from app.agents.entity_extractor import EntityExtractorAgent +from app.agents.orchestrator import AgentOrchestrator +from app.agents.safety_analyzer import SafetyAnalyzerAgent +from app.agents.summarizer import SummarizerAgent + +__all__ = [ + "BaseAgent", + "ClassifierAgent", + "SafetyAnalyzerAgent", + "EntityExtractorAgent", + "SummarizerAgent", + "AgentOrchestrator", +] diff --git a/app/agents/base.py b/app/agents/base.py new file mode 100644 index 0000000000000000000000000000000000000000..6972f69a2e702b8a6e7809ab34c3f59fd7d838b1 --- /dev/null +++ b/app/agents/base.py @@ -0,0 +1,308 @@ +""" +Base Agent +Abstract base class for all mining intelligence agents. + +Improvements over v1: + - JSON mode via response_mime_type (no more regex JSON extraction) + - Retry with exponential backoff (3 retries), respecting Gemini retry_delay + - Processes full document via pages, not truncated to 3000 chars + - Confidence score required in all agent outputs + - Proper QuotaExceededError raised (no more silent empty-dict returns) +""" + +import asyncio +import logging +import re +from abc import ABC, abstractmethod +from typing import Any, Dict, List, Optional + +import google.generativeai as genai +from google.generativeai.types import GenerationConfig + +from app.config import settings + +logger = logging.getLogger(__name__) + +# Configure Gemini once at module level +genai.configure(api_key=settings.GEMINI_API_KEY) + +# Generation config that forces JSON output — no more regex parsing +_JSON_GENERATION_CONFIG = GenerationConfig( + response_mime_type="application/json", + temperature=0.1, # Low temperature for consistent structured output + top_p=0.95, +) + +_MAX_RETRIES = 3 +_RETRY_BASE_DELAY = 2.0 # seconds — minimum delay between retries +_MAX_RETRY_DELAY = 120.0 # seconds — cap for retry_delay parsed from API response + + +class QuotaExceededError(RuntimeError): + """Raised when the Gemini API quota / rate-limit is exhausted.""" + + +def _parse_retry_delay(err_str: str) -> Optional[float]: + """ + Extract the suggested retry_delay (in seconds) from a Gemini 429 error + message. The error body contains a line like: + retry_delay { seconds: 31 } + Returns None if no delay can be parsed. + """ + match = re.search(r"retry_delay\s*\{\s*seconds:\s*(\d+)", err_str) + if match: + return min(float(match.group(1)), _MAX_RETRY_DELAY) + # Fallback: look for "Please retry in X.Xs" + match2 = re.search(r"retry in (\d+\.?\d*)s", err_str) + if match2: + return min(float(match2.group(1)), _MAX_RETRY_DELAY) + return None + + +class BaseAgent(ABC): + """ + Abstract base class for mining document intelligence agents. + + Each agent is responsible for a specific analysis task: + - Classification + - Safety Analysis + - Entity Extraction + - Summarization + """ + + def __init__( + self, + model_name: str = None, + provider: str = "gemini", + fallback_model: str = None, + fallback_provider: str = None, + ): + self.provider = provider + self.model_name = model_name or settings.GEMINI_MODEL + self.name = self.__class__.__name__ + + # Fallback config (e.g. Cerebras when Groq is rate-limited) + self.fallback_model = fallback_model + self.fallback_provider = fallback_provider + self._fallback_client = None + self._using_fallback = False + + if self.fallback_provider and self.fallback_model: + self._init_fallback_client() + + self._init_client() + + def _init_client(self): + """Initialize the primary provider client.""" + if self.provider == "gemini": + self.model = genai.GenerativeModel( + model_name=self.model_name, + generation_config=_JSON_GENERATION_CONFIG, + ) + elif self.provider == "groq": + from app.services.llm_provider import get_groq_client + + self.client = get_groq_client() + elif self.provider == "mistral": + from app.services.llm_provider import get_mistral_client + + self.client = get_mistral_client() + elif self.provider == "cerebras": + from app.services.llm_provider import get_cerebras_client + + self.client = get_cerebras_client() + + def _init_fallback_client(self): + """Initialize the fallback provider client.""" + if self.fallback_provider == "cerebras": + from app.services.llm_provider import get_cerebras_client + + self._fallback_client = get_cerebras_client() + elif self.fallback_provider == "groq": + from app.services.llm_provider import get_groq_client + + self._fallback_client = get_groq_client() + elif self.fallback_provider == "mistral": + from app.services.llm_provider import get_mistral_client + + self._fallback_client = get_mistral_client() + logger.info( + f"{self.name}: Fallback configured — {self.fallback_provider}/{self.fallback_model}" + ) + + @abstractmethod + async def analyze( + self, + text: str, + context: Optional[Dict] = None, + ) -> Dict[str, Any]: + """ + Analyze document text and return structured results. + + Args: + text: Document text content (full text or representative sample) + context: Additional context (e.g., document category from classifier) + + Returns: + Dictionary with agent-specific analysis results. + All results MUST include a 'confidence' key (0.0–1.0). + """ + + @property + @abstractmethod + def system_prompt(self) -> str: + """System prompt defining the agent's role and capabilities.""" + + # ── Generation with retry + fallback ─────────────────────────────────────── + + async def _call_openai_compat(self, client, model: str, prompt: str) -> str: + """Call an OpenAI-compatible provider and return text response.""" + response = await client.chat.completions.create( + model=model, + messages=[ + {"role": "system", "content": self.system_prompt}, + {"role": "user", "content": prompt}, + ], + response_format={"type": "json_object"}, + temperature=0.1, + ) + return response.choices[0].message.content or "" + + async def _generate_json(self, prompt: str) -> Dict[str, Any]: + """ + Generate structured JSON output with retry and automatic provider fallback. + + Primary provider is tried first. On 429/rate-limit errors, if a fallback + is configured (e.g. Cerebras when Groq is rate-limited), the request + is retried on the fallback provider before raising QuotaExceededError. + """ + import json + + last_error: Optional[Exception] = None + for attempt in range(1, _MAX_RETRIES + 1): + try: + full_prompt = f"{self.system_prompt}\n\n{prompt}" + + if self.provider == "gemini": + response = await asyncio.to_thread( + self.model.generate_content, + full_prompt, + ) + text_response = getattr(response, "text", "") + elif self.provider in ["groq", "mistral", "cerebras"]: + text_response = await self._call_openai_compat( + self.client, self.model_name, prompt + ) + + try: + return json.loads(text_response) + except (json.JSONDecodeError, AttributeError) as parse_err: + logger.warning( + f"{self.name} attempt {attempt}: JSON parse failed — {parse_err}. " + f"Raw response: {text_response[:200]}" + ) + last_error = parse_err + delay = _RETRY_BASE_DELAY * (2 ** (attempt - 1)) + await asyncio.sleep(delay) + continue + + except Exception as e: + last_error = e + err_str = str(e) + + is_quota = ( + "429" in err_str + or "rate_limit" in err_str.lower() + or "quota" in err_str.lower() + or "RESOURCE_EXHAUSTED" in err_str + ) + + if is_quota: + # Try fallback provider if available and not already using it + if ( + self._fallback_client + and self.fallback_model + and not self._using_fallback + ): + logger.warning( + f"{self.name}: Primary provider rate-limited. " + f"Falling back to {self.fallback_provider}/{self.fallback_model}" + ) + try: + text_response = await self._call_openai_compat( + self._fallback_client, self.fallback_model, prompt + ) + result = json.loads(text_response) + self._using_fallback = True + return result + except Exception as fb_err: + logger.error(f"{self.name}: Fallback also failed: {fb_err}") + # Fall through to raise QuotaExceededError + + # Parse the suggested wait time from the error body + suggested_delay = _parse_retry_delay(err_str) + + if attempt < _MAX_RETRIES and suggested_delay is not None: + logger.warning( + f"{self.name}: Quota/rate-limit hit (attempt {attempt}/{_MAX_RETRIES}). " + f"Waiting {suggested_delay}s..." + ) + await asyncio.sleep(suggested_delay) + continue + + logger.error(f"{self.name}: All providers exhausted — {e}") + raise QuotaExceededError( + f"Rate limit exceeded for {self.name}. " + "Please try again later." + ) from e + + # Transient non-quota error — exponential backoff + delay = _RETRY_BASE_DELAY * (2 ** (attempt - 1)) + logger.warning( + f"{self.name} attempt {attempt}/{_MAX_RETRIES} failed: {e}. " + f"Retrying in {delay}s..." + ) + await asyncio.sleep(delay) + + logger.error(f"{self.name} failed after {_MAX_RETRIES} attempts: {last_error}") + return {} + + # ── Text helpers ─────────────────────────────────────────────────────────── + + def _prepare_text(self, text: str, max_chars: int = 15000) -> str: + """ + Prepare text for agent analysis. + + Instead of hard-truncating to 3000 chars (old behavior), we use up to + 15000 chars (≈10 pages) to capture much more document content. + Long documents get the first 12000 chars + last 3000 chars to include + both the opening context and the conclusion/summary sections. + """ + if len(text) <= max_chars: + return text + + head = text[:12000] + tail = text[-3000:] + return head + "\n\n[... middle of document omitted for analysis ...]\n\n" + tail + + # ── Kept for backward compatibility ─────────────────────────────────────── + + def _parse_json(self, text: str) -> Dict[str, Any]: + """Legacy JSON parser — kept for any subclass that still needs it.""" + import json + import re + + text = text.strip() + if text.startswith("```"): + lines = text.split("\n") + text = "\n".join(line for line in lines if not line.startswith("```")) + try: + return json.loads(text) + except json.JSONDecodeError: + match = re.search(r"\{.*\}", text, re.DOTALL) + if match: + try: + return json.loads(match.group()) + except json.JSONDecodeError: + pass + return {} diff --git a/app/agents/classifier.py b/app/agents/classifier.py new file mode 100644 index 0000000000000000000000000000000000000000..92b053cc8a68b76b32fd0bea6efa3a5d85396b7d --- /dev/null +++ b/app/agents/classifier.py @@ -0,0 +1,107 @@ +""" +Classifier Agent +Document classification for mining industry categories +""" + +import asyncio +import logging +from typing import Any, Dict, Optional + +from app.agents.base import BaseAgent +from app.models.document import DocumentCategory + +logger = logging.getLogger(__name__) + + +class ClassifierAgent(BaseAgent): + """ + Document Classification Agent. + + Categorizes mining documents into predefined categories: + - Safety protocols + - Equipment manuals + - Regulatory documents + - Incident reports + - Geological reports + - Environmental reports + - Training materials + - Permits + - Maintenance logs + """ + + def __init__(self): + super().__init__(model_name="llama-3.3-70b-versatile", provider="groq") + + @property + def system_prompt(self) -> str: + return """You are a document classification agent specialized in the mining industry. + +Your task is to analyze documents and classify them into the appropriate category based on their content, structure, and purpose. + +Categories: +1. safety_protocol - Safety procedures, guidelines, emergency protocols +2. equipment_manual - Equipment operation guides, maintenance manuals +3. regulatory - MSHA, OSHA, EPA regulations, compliance documents +4. incident_report - Accident reports, incident investigations, near-miss reports +5. geological - Drill logs, assay reports, geological surveys, core samples +6. environmental - Environmental impact assessments, monitoring reports +7. training - Training materials, certifications, competency assessments +8. permit - Mining permits, licenses, applications +9. maintenance - Maintenance schedules, repair logs, equipment inspections +10. other - Documents that don't fit other categories + +Consider: +- Document structure and formatting +- Key terminology and language used +- Purpose and intended audience +- Regulatory references +""" + + async def analyze( + self, text: str, context: Optional[Dict] = None + ) -> Dict[str, Any]: + """ + Classify document into mining category. + + Returns: + { + "category": str (DocumentCategory value), + "subcategory": str, + "confidence": float (0-1), + "reasoning": str + } + """ + prompt = f"""Analyze this mining document and classify it. + +Document content ({len(text)} chars total, showing up to 15000): +{self._prepare_text(text)} + +Respond with a JSON object: +{{ + "category": "", + "subcategory": "", + "confidence": <0.0-1.0>, + "reasoning": "" +}} +""" + result = await self._generate_json(prompt) + + category_str = (result.get("category") or "other").lower().strip() + category_map = { + "safety_protocol": DocumentCategory.SAFETY_PROTOCOL, + "equipment_manual": DocumentCategory.EQUIPMENT_MANUAL, + "regulatory": DocumentCategory.REGULATORY, + "incident_report": DocumentCategory.INCIDENT_REPORT, + "geological": DocumentCategory.GEOLOGICAL, + "environmental": DocumentCategory.ENVIRONMENTAL, + "training": DocumentCategory.TRAINING, + "permit": DocumentCategory.PERMIT, + "maintenance": DocumentCategory.MAINTENANCE, + } + + return { + "category": category_map.get(category_str, DocumentCategory.OTHER).value, + "subcategory": result.get("subcategory"), + "confidence": float(result.get("confidence") or 0.5), + "reasoning": result.get("reasoning", ""), + } diff --git a/app/agents/compliance_auditor.py b/app/agents/compliance_auditor.py new file mode 100644 index 0000000000000000000000000000000000000000..cbdb10a34c7b805701acde0721ffd6a7b6615ad6 --- /dev/null +++ b/app/agents/compliance_auditor.py @@ -0,0 +1,129 @@ +""" +Compliance Auditor Agent +Cross-references a regulation clause against operational document evidence +to determine compliance status (compliant / gap / missing). + +Uses Gemini for the nuanced cross-referencing task that requires reasoning +across multiple evidence chunks. +""" + +import logging +from typing import Any, Dict, List, Optional + +from app.agents.base import BaseAgent + +logger = logging.getLogger(__name__) + + +class ComplianceAuditorAgent(BaseAgent): + """ + Regulatory Compliance Auditor Agent. + + Takes a single regulation clause and a set of evidence chunks from + operational documents, then assesses whether the operational documents + adequately address the clause requirements. + """ + + def __init__(self): + super().__init__(model_name="llama-3.3-70b-versatile", provider="groq") + + @property + def system_prompt(self) -> str: + return """You are a regulatory compliance auditor specializing in the mining industry. + +Your expertise includes: +- MSHA (Mine Safety and Health Administration) regulations (30 CFR) +- OSHA safety standards (29 CFR 1910, 1926) +- DGMS (Directorate General of Mines Safety) regulations +- EPA environmental regulations for mining operations +- State-level mining regulations and permits + +Your task: Given a REGULATION CLAUSE and EVIDENCE CHUNKS from operational documents, +assess whether the operational documents adequately address the clause requirements. + +Assessment statuses: +- "compliant": The operational documents clearly address the regulation clause requirements +- "gap": The operational documents partially address the clause but have gaps or deficiencies +- "missing": The operational documents do not address this clause at all, or the evidence is insufficient + +Be strict but fair. If the evidence is thin but directionally correct, mark as "gap" not "missing". +Only mark "compliant" if the evidence clearly and adequately addresses the clause. + +Always cite specific evidence in your assessment. If no evidence is provided, mark as "missing". +""" + + async def analyze( + self, + text: str, + context: Optional[Dict] = None, + ) -> Dict[str, Any]: + """ + Assess compliance for a single regulation clause. + + Args: + text: The regulation clause text + context: Must contain: + - evidence_chunks: List of dicts with chunk_text, document_title, + page_numbers, relevance_score + - clause_section: Section title from the regulation document + + Returns: + { + "status": "compliant" | "gap" | "missing", + "assessment": str, + "confidence": float (0.0-1.0), + "recommendations": list[str], + } + """ + evidence_chunks: List[Dict] = ( + context.get("evidence_chunks", []) if context else [] + ) + clause_section = context.get("clause_section", "") if context else "" + + # Format evidence for the prompt + if evidence_chunks: + evidence_text = "\n\n".join( + f"[Evidence {i+1}] From '{chunk.get('document_title', 'Unknown')}'" + f" (Pages {chunk.get('page_numbers', ['?'])}, " + f"Relevance: {chunk.get('relevance_score', 0):.0%}):\n" + f"{chunk.get('chunk_text', '')}" + for i, chunk in enumerate(evidence_chunks) + ) + else: + evidence_text = "No relevant evidence found in operational documents." + + section_hint = ( + f"\nRegulation Section: {clause_section}" if clause_section else "" + ) + + prompt = ( + "Assess compliance for the following regulation clause against " + "the provided operational document evidence.\n\n" + f"REGULATION CLAUSE{text}:\n{text}\n\n" + f"{section_hint}\n\n" + f"EVIDENCE FROM OPERATIONAL DOCUMENTS:\n{evidence_text}\n\n" + "Respond with a JSON object:\n" + "{\n" + ' "status": "",\n' + ' "assessment": "<2-4 sentence explanation of compliance status, ' + 'citing specific evidence where available>",\n' + ' "confidence": <0.0-1.0>,\n' + ' "recommendations": [""]\n' + "}\n" + ) + + result = await self._generate_json(prompt) + + status = (result.get("status") or "missing").lower() + if status not in ("compliant", "gap", "missing"): + status = "missing" + + return { + "status": status, + "assessment": result.get( + "assessment", + "Assessment could not be generated.", + ), + "confidence": float(result.get("confidence") or 0.5), + "recommendations": result.get("recommendations", []), + } diff --git a/app/agents/entity_extractor.py b/app/agents/entity_extractor.py new file mode 100644 index 0000000000000000000000000000000000000000..23d973ea1026392e09ecbc24b6763b8afe351f87 --- /dev/null +++ b/app/agents/entity_extractor.py @@ -0,0 +1,130 @@ +""" +Entity Extractor Agent +Mining-specific Named Entity Recognition +""" + +import logging +from typing import Any, Dict, List, Optional + +from app.agents.base import BaseAgent + +logger = logging.getLogger(__name__) + + +class EntityExtractorAgent(BaseAgent): + """ + Named Entity Recognition Agent for Mining Documents. + + Extracts mining-specific entities: + - Equipment names and models + - Chemical compounds and gases + - Mine locations and sections + - Personnel and roles + - Dates and schedules + - Regulatory references + """ + + def __init__(self): + # Cerebras: 1M tokens/day free, 2600+ TPS, 60K TPM — better for high-volume extraction than Groq + super().__init__(model_name="gpt-oss-120b", provider="cerebras") + + @property + def system_prompt(self) -> str: + return """You are a named entity extraction agent specialized in mining documents. + +Extract the following entity types: + +1. EQUIPMENT + - Mining machinery (excavators, haul trucks, drills) + - Brand names and models (Caterpillar D11, Komatsu PC8000) + - Equipment IDs and serial numbers + - Tools and instruments + +2. CHEMICALS + - Gases (methane, CO, H2S, oxygen) + - Minerals and ores + - Explosives and blasting agents + - Dust types (coal dust, silica) + - Hazardous substances + +3. LOCATIONS + - Mine names + - Sections and portals + - Underground levels + - Surface areas + - Geographic coordinates + +4. PERSONNEL + - Names (anonymize if needed) + - Roles (Safety Officer, Foreman, Engineer) + - Departments and teams + - Certifications + +5. DATES + - Specific dates + - Deadlines + - Scheduled events + - Time periods + +6. REGULATIONS + - MSHA regulations (30 CFR citations) + - OSHA standards + - EPA requirements + - State regulations + - Company policies + +Be precise and avoid duplicates. Extract exactly as written in the document. +""" + + async def analyze( + self, text: str, context: Optional[Dict] = None + ) -> Dict[str, Any]: + """ + Extract named entities from document. + + Returns dict with keys: equipment, chemicals, locations, personnel, + dates, regulations (all List[str]), entity_count (int) + """ + prompt = f"""Extract all mining-specific named entities from this document. + +Document content ({len(text)} chars total, showing up to 15000): +{self._prepare_text(text)} + +Respond with a JSON object: +{{ + "equipment": [""], + "chemicals": [""], + "locations": [""], + "personnel": [""], + "dates": [""], + "regulations": [""] +}} + +Notes: +- List each unique entity only once +- Use exact text from document +- For personnel, prefer roles over names for privacy +- Include regulation citations in standard format +""" + result = await self._generate_json(prompt) + + entities = { + "equipment": self._deduplicate(result.get("equipment", [])), + "chemicals": self._deduplicate(result.get("chemicals", [])), + "locations": self._deduplicate(result.get("locations", [])), + "personnel": self._deduplicate(result.get("personnel", [])), + "dates": self._deduplicate(result.get("dates", [])), + "regulations": self._deduplicate(result.get("regulations", [])), + } + entities["entity_count"] = sum(len(v) for v in entities.values()) + return entities + + def _deduplicate(self, items: List[str]) -> List[str]: + """Remove duplicates while preserving order""" + seen = set() + result = [] + for item in items: + if item and item.lower() not in seen: + seen.add(item.lower()) + result.append(item) + return result diff --git a/app/agents/orchestrator.py b/app/agents/orchestrator.py new file mode 100644 index 0000000000000000000000000000000000000000..e8c5ac155a7f97b834db94c035266f0c2307b0cd --- /dev/null +++ b/app/agents/orchestrator.py @@ -0,0 +1,252 @@ +""" +Agent Orchestrator +Coordinates multi-agent document analysis pipeline. + +Improvements over v1: + - Accepts pages parameter from extractor for future per-page analysis + - Runs safety/entity/summary with small delay between calls to avoid simultaneous quota hits + - Adds per-agent timing metrics + - All agents use JSON mode (no regex) + exponential backoff retry + - Surfaces QuotaExceededError instead of silently returning empty data +""" + +import asyncio +import logging +from datetime import datetime, timezone +from typing import Any, Dict, List, Optional + +from app.agents.base import QuotaExceededError +from app.agents.classifier import ClassifierAgent +from app.agents.entity_extractor import EntityExtractorAgent +from app.agents.safety_analyzer import SafetyAnalyzerAgent +from app.agents.summarizer import SummarizerAgent + +logger = logging.getLogger(__name__) + + +import functools + + +def exponential_backoff_wrapper(max_retries: int = 3, base_delay: float = 2.0): + """Exponential Backoff Utility Wrapper for AI Agent execution pipelines.""" + + def decorator(func): + @functools.wraps(func) + async def wrapper(*args, **kwargs): + last_error = None + for attempt in range(1, max_retries + 1): + try: + return await func(*args, **kwargs) + except Exception as e: + err_str = str(e).lower() + is_rate_limit = any( + term in err_str + for term in [ + "429", + "quota", + "timeout", + "rate_limit", + "resource_exhausted", + "too many requests", + ] + ) + if is_rate_limit: + last_error = e + if attempt < max_retries: + delay = base_delay * (2 ** (attempt - 1)) + logger.warning( + f"Audit Trail: {func.__name__} attempt {attempt} failed (429/Timeout). Retrying in {delay}s..." + ) + await asyncio.sleep(delay) + else: + logger.error( + f"Audit Trail: {func.__name__} exhausted {max_retries} retries." + ) + raise QuotaExceededError( + f"Rate limit/timeout exceeded in {func.__name__}" + ) from e + else: + raise + if last_error is not None: + raise last_error + raise RuntimeError(f"Task {func.__name__} failed with no attempts made") + + return wrapper + + return decorator + + +class AgentOrchestrator: + """ + Multi-Agent Orchestrator for Document Intelligence. + + Execution order: + 1. ClassifierAgent — runs first (result feeds category context to others) + 2. SafetyAnalyzerAgent ┐ + 3. EntityExtractorAgent ├── run in parallel after classification (Multi-Provider) + 4. SummarizerAgent ┘ + """ + + def __init__(self): + self.classifier = ClassifierAgent() + self.safety_analyzer = SafetyAnalyzerAgent() + self.entity_extractor = EntityExtractorAgent() + self.summarizer = SummarizerAgent() + + async def analyze_document( + self, + text: str, + pages: Optional[List] = None, + ) -> Dict[str, Any]: + """ + Run full multi-agent analysis pipeline on document. + """ + start_time = datetime.now(timezone.utc) + logger.info("Starting multi-agent document analysis") + + agent_timings: Dict[str, int] = {} + + try: + # ── Step 1: Classification (feeds category context to other agents) ── + t0 = datetime.now(timezone.utc) + logger.info("Running ClassifierAgent...") + + @exponential_backoff_wrapper() + async def _run_classifier(): + return await self.classifier.analyze(text) + + classification = await _run_classifier() + agent_timings["classifier_ms"] = int( + (datetime.now(timezone.utc) - t0).total_seconds() * 1000 + ) + + category = classification.get("category", "other") + context = {"category": category} + + # ── Step 2: Parallel agents using category context ───────────────── + logger.info( + "Running parallel agents across multiple providers (Safety, Entity, Summary)..." + ) + t1 = datetime.now(timezone.utc) + + @exponential_backoff_wrapper() + async def _run_safety(): + # Only run safety analysis on relevant document categories + non_safety_categories = [ + "regulatory", + "geological", + "environmental", + "permit", + "other", + ] + if category in non_safety_categories: + logger.info( + f"Routing: Document is {category}. Bypassing Safety Analyzer." + ) + return { + "status": "not_applicable", + "score": None, + "hazards": [], + "recommendations": [ + f"Safety analysis bypassed for {category} document" + ], + } + return await self.safety_analyzer.analyze(text, context) + + @exponential_backoff_wrapper() + async def _run_entities(): + return await self.entity_extractor.analyze(text, context) + + @exponential_backoff_wrapper() + async def _run_summary(): + return await self.summarizer.analyze(text, context) + + safety, entities, summary = await asyncio.gather( + _run_safety(), + _run_entities(), + _run_summary(), + return_exceptions=True, + ) + + agent_timings["parallel_agents_ms"] = int( + (datetime.now(timezone.utc) - t1).total_seconds() * 1000 + ) + + # Handle per-agent exceptions gracefully + def _quota_error_result(agent_name: str, exc: Exception) -> dict: + is_quota = isinstance(exc, QuotaExceededError) + logger.error(f"{agent_name} failed: {exc}") + return { + "error": str(exc), + "quota_exceeded": is_quota, + "status": "quota_exceeded" if is_quota else "error", + } + + if isinstance(safety, Exception): + safety = { + **_quota_error_result("SafetyAnalyzerAgent", safety), + "score": None, + "hazards": [], + "recommendations": [], + } + + if isinstance(entities, Exception): + entities = { + **_quota_error_result("EntityExtractorAgent", entities), + "equipment": [], + "chemicals": [], + "locations": [], + "personnel": [], + "dates": [], + "regulations": [], + } + + if isinstance(summary, Exception): + summary = { + **_quota_error_result("SummarizerAgent", summary), + "summary": "Analysis failed — Gemini API quota exceeded. Please try again later.", + "key_points": [], + } + + total_ms = int( + (datetime.now(timezone.utc) - start_time).total_seconds() * 1000 + ) + logger.info(f"Multi-agent analysis completed in {total_ms}ms") + + return { + "classification": classification, + "safety": safety, + "entities": entities, + "summary": summary, + "metadata": { + "processing_time_ms": total_ms, + "agent_timings": agent_timings, + "agents_used": [ + "classifier", + "safety_analyzer", + "entity_extractor", + "summarizer", + ], + "analyzed_at": datetime.now(timezone.utc).isoformat(), + }, + } + + except QuotaExceededError: + # Re-raise quota errors so document_service.py can handle them + # with its dedicated QuotaExceededError handler (partial save, not FAILED). + raise + + except Exception as e: + logger.error(f"Orchestrator failed: {e}", exc_info=True) + return { + "error": str(e), + "metadata": {"failed": True, "error_message": str(e)}, + } + + async def analyze_for_safety_only(self, text: str) -> Dict[str, Any]: + """Quick safety-only analysis for real-time checks.""" + return await self.safety_analyzer.analyze(text) + + async def classify_only(self, text: str) -> Dict[str, Any]: + """Quick classification only.""" + return await self.classifier.analyze(text) diff --git a/app/agents/safety_analyzer.py b/app/agents/safety_analyzer.py new file mode 100644 index 0000000000000000000000000000000000000000..6f0ee765c8cc76c89dc0defa6920ffdef3ea6241 --- /dev/null +++ b/app/agents/safety_analyzer.py @@ -0,0 +1,140 @@ +""" +Safety Analyzer Agent +Mining safety compliance and hazard detection. + +Uses JSON mode (response_mime_type=application/json) + retry for reliable output. +Processes up to 15000 chars of document content (vs 5000 in v1). +""" + +import logging +from typing import Any, Dict, List, Optional + +from app.agents.base import BaseAgent + +logger = logging.getLogger(__name__) + + +class SafetyAnalyzerAgent(BaseAgent): + """ + Safety Compliance Analysis Agent. + + Analyzes mining documents for: + - MSHA/OSHA/DGMS compliance issues + - Safety hazards and risks + - Missing safety requirements + - Recommendations for improvement + """ + + def __init__(self): + # magistral-small-latest is a reasoning model, not a standard chat model. + # It uses chain-of-thought internally for compliance scoring. + super().__init__(model_name="magistral-small-latest", provider="mistral") + + @property + def system_prompt(self) -> str: + return """You are a mining safety compliance analysis agent. + +Your expertise includes: +- MSHA (Mine Safety and Health Administration) regulations +- OSHA safety standards +- DGMS (Directorate General of Mines Safety) regulations +- Underground and surface mining safety +- Equipment safety requirements +- Emergency response protocols +- Ventilation and air quality standards +- Ground control and stability +- Electrical safety in mining +- Personal protective equipment (PPE) +- Hazardous materials handling + +Analyze documents for: +1. Compliance with regulations +2. Potential safety hazards +3. Missing safety procedures +4. Risk factors +5. Areas needing improvement + +Be thorough but practical. Focus on actionable findings. +""" + + async def analyze( + self, text: str, context: Optional[Dict] = None + ) -> Dict[str, Any]: + """ + Analyze document for safety compliance and hazards. + + Returns: + { + "score": float (0-100), + "status": str ("compliant"|"warning"|"violation"), + "hazards": List[Dict], + "recommendations": List[str], + "compliance_details": Dict, + "confidence": float (0-1), + "reasoning": Dict (explainability layer) + } + """ + category = context.get("category", "unknown") if context else "unknown" + + prompt = ( + "Analyze this mining document for safety compliance and hazards.\n\n" + f"Document Type: {category}\n" + f"Document content ({len(text)} chars total, showing up to 15000):\n" + f"{self._prepare_text(text)}\n\n" + "Evaluate and respond with a JSON object:\n" + "{\n" + ' "score": <0-100 overall safety score>,\n' + ' "status": "",\n' + ' "confidence": <0.0-1.0>,\n' + ' "hazards": [\n' + " {\n" + ' "type": "",\n' + ' "severity": "",\n' + ' "description": "",\n' + ' "regulation": ""\n' + " }\n" + " ],\n" + ' "recommendations": [""],\n' + ' "compliance_details": {\n' + ' "msha_compliant": ,\n' + ' "osha_compliant": ,\n' + ' "dgms_compliant": ,\n' + ' "missing_elements": [""]\n' + " },\n" + ' "reasoning": {\n' + ' "score_explanation": "<1-2 sentence explanation of why this specific score was assigned>",\n' + ' "positive_factors": [""],\n' + ' "negative_factors": [""],\n' + ' "evidence": [\n' + ' {"text": "", "factor": "", "impact": ""}\n' + " ]\n" + " },\n" + ' "summary": ""\n' + "}\n\n" + "Scoring guide:\n" + " 80-100: Compliant, minimal concerns\n" + " 60-79: Generally compliant with warnings\n" + " 40-59: Significant concerns, needs attention\n" + " 0-39: Critical issues or violations present\n" + ) + + result = await self._generate_json(prompt) + + return { + "score": float(result.get("score") or 50), + "status": (result.get("status") or "pending").lower(), + "confidence": float(result.get("confidence") or 0.5), + "hazards": result.get("hazards", []), + "recommendations": result.get("recommendations", []), + "compliance_details": result.get("compliance_details", {}), + "reasoning": result.get( + "reasoning", + { + "score_explanation": "", + "positive_factors": [], + "negative_factors": [], + "evidence": [], + }, + ), + "summary": result.get("summary", ""), + } diff --git a/app/agents/summarizer.py b/app/agents/summarizer.py new file mode 100644 index 0000000000000000000000000000000000000000..db0d9347aca626fb1fe61efefaf33635f17dc65b --- /dev/null +++ b/app/agents/summarizer.py @@ -0,0 +1,100 @@ +""" +Summarizer Agent +Document summarization for mining documents +""" + +import logging +from typing import Any, Dict, List, Optional + +from app.agents.base import BaseAgent + +logger = logging.getLogger(__name__) + + +class SummarizerAgent(BaseAgent): + """ + Document Summarization Agent. + + Creates concise, actionable summaries of mining documents: + - Executive summary + - Key points extraction + - Action items identification + """ + + def __init__(self): + # Shift to Cerebras (gpt-oss-120b) to bypass Gemini API rate limits and avoid Groq parallel execution rate limits + super().__init__(model_name="gpt-oss-120b", provider="cerebras") + + @property + def system_prompt(self) -> str: + return """You are a document summarization agent for the mining industry. + +Create clear, actionable summaries that: +- Highlight critical information first +- Focus on safety-relevant content +- Identify action items and deadlines +- Use plain language accessible to all mining personnel +- Preserve technical accuracy + +Summary structure: +1. Executive Summary: 2-3 paragraphs covering main purpose and findings +2. Key Points: 5-7 bullet points of most important information +3. Action Items: Any required actions or follow-ups (if applicable) + +Prioritize: +- Safety information +- Compliance requirements +- Deadlines and schedules +- Equipment status +- Personnel responsibilities +""" + + async def analyze( + self, text: str, context: Optional[Dict] = None + ) -> Dict[str, Any]: + """ + Generate document summary and key points. + + Returns: + { + "summary": str, + "key_points": List[str], + "action_items": List[str], + "document_purpose": str, + "confidence": float + } + """ + category = context.get("category", "unknown") if context else "unknown" + + prompt = f"""Summarize this mining document clearly and concisely. + +Document Type: {category} +Document content ({len(text)} chars total, showing up to 15000): +{self._prepare_text(text)} + +Respond with a JSON object: +{{ + "summary": "<2-3 paragraph executive summary covering main purpose and findings>", + "key_points": [ + "", + "", + "", + "", + "" + ], + "action_items": [""], + "document_purpose": "", + "confidence": <0.0-1.0> +}} +""" + result = await self._generate_json(prompt) + + summary = result.get("summary") or "Summary not available." + return { + "summary": summary, + "key_points": result.get("key_points", []), + "action_items": result.get("action_items", []), + "document_purpose": result.get("document_purpose", ""), + "confidence": float(result.get("confidence") or 0.7), + "word_count": len(summary.split()), + } diff --git a/app/api/__init__.py b/app/api/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..db18ae4becb3645c570db4ba49ffb65a27ab6a0c --- /dev/null +++ b/app/api/__init__.py @@ -0,0 +1,4 @@ +""" +API Module +REST API endpoints organized by domain +""" diff --git a/app/api/deps.py b/app/api/deps.py new file mode 100644 index 0000000000000000000000000000000000000000..5c080c13360ecb6a15112aa8ac0711bae1779e2c --- /dev/null +++ b/app/api/deps.py @@ -0,0 +1,112 @@ +""" +API Dependencies +Shared dependencies for FastAPI endpoints +""" + +import hashlib +import logging +from typing import Optional + +from fastapi import Depends, Header, Request +from fastapi.security import HTTPAuthorizationCredentials, HTTPBearer +from sqlalchemy.orm import Session + +from app.core.exceptions import AuthenticationError +from app.core.security import extract_user_email, extract_user_id, verify_jwt_token +from app.db.session import SessionLocal, get_db +from app.models.audit import AuditAction, create_audit_log +from app.models.user import User + +logger = logging.getLogger(__name__) + + +def _anonymize_user_id(user_id: str) -> str: + """Create a truncated hash of user_id for safe logging.""" + return hashlib.sha256(user_id.encode()).hexdigest()[:12] + + +# Security scheme +security = HTTPBearer() + + +async def get_current_user_id( + credentials: HTTPAuthorizationCredentials = Depends(security), +) -> str: + """ + Dependency to extract and verify user ID from JWT. + Returns Clerk user ID string. + """ + token = credentials.credentials + payload = await verify_jwt_token(token) + return extract_user_id(payload) + + +async def get_current_user( + user_id: str = Depends(get_current_user_id), db: Session = Depends(get_db) +) -> User: + """ + Dependency to get current user model from database. + Creates user record if it doesn't exist (first login). + """ + user = db.query(User).filter(User.clerk_user_id == user_id).first() + + if not user: + # Auto-create user on first access + user = User(clerk_user_id=user_id, is_active=True) + db.add(user) + db.commit() + db.refresh(user) + logger.info(f"Created new user: {_anonymize_user_id(user_id)}") + + return user + + +async def get_optional_user( + credentials: Optional[HTTPAuthorizationCredentials] = Depends( + HTTPBearer(auto_error=False) + ), + db: Session = Depends(get_db), +) -> Optional[User]: + """ + Dependency for endpoints that work with or without auth. + Returns User if authenticated, None otherwise. + """ + if not credentials: + return None + + try: + token = credentials.credentials + payload = await verify_jwt_token(token) + user_id = extract_user_id(payload) + return db.query(User).filter(User.clerk_user_id == user_id).first() + except Exception: + return None + + +def get_client_ip(request: Request) -> str: + """Extract client IP address from request""" + forwarded = request.headers.get("X-Forwarded-For") + if forwarded: + return forwarded.split(",")[0].strip() + return request.client.host if request.client else "unknown" + + +def get_user_agent(request: Request) -> str: + """Extract user agent from request""" + return request.headers.get("User-Agent", "unknown") + + +async def audit_middleware( + request: Request, user_id: str = Depends(get_current_user_id) +): + """ + Middleware-like dependency to log API access. + Add to endpoints that need audit logging. + """ + # This is called after auth, so we have user_id + # Actual logging happens in endpoint handlers + return { + "user_id": user_id, + "ip_address": get_client_ip(request), + "user_agent": get_user_agent(request), + } diff --git a/app/api/v1/__init__.py b/app/api/v1/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..578e3de85c425533d191eb0a7243a73c2d97ebcb --- /dev/null +++ b/app/api/v1/__init__.py @@ -0,0 +1,8 @@ +""" +API v1 Module +Versioned API endpoints +""" + +from app.api.v1.router import api_router + +__all__ = ["api_router"] diff --git a/app/api/v1/analytics.py b/app/api/v1/analytics.py new file mode 100644 index 0000000000000000000000000000000000000000..fb7ff743474ad81902ea0d12309b81416d80b622 --- /dev/null +++ b/app/api/v1/analytics.py @@ -0,0 +1,410 @@ +""" +Analytics API Endpoints +Dashboard statistics and mining intelligence metrics +""" + +import logging +from datetime import datetime, timedelta + +from fastapi import APIRouter, Depends, Query +from sqlalchemy import case, func +from sqlalchemy.orm import Session + +from app.api.deps import get_current_user_id +from app.db.session import get_db +from app.models.chat import ChatMessage, ChatSession +from app.models.document import ( + ComplianceStatus, + Document, + DocumentCategory, + DocumentStatus, +) +from app.schemas.analytics import ( + CategoryCount, + DashboardStats, + DocumentAnalytics, + SafetyAnalytics, + SafetyDistribution, + StatusCount, +) + +logger = logging.getLogger(__name__) + +router = APIRouter() + + +import json + +import redis + +from app.config import settings + +# Optional Redis client +try: + redis_client = redis.Redis.from_url( + settings.REDIS_URL, decode_responses=True, socket_timeout=1 + ) +except Exception: + redis_client = None + + +@router.get("/dashboard", response_model=DashboardStats) +async def get_dashboard_stats( + user_id: str = Depends(get_current_user_id), db: Session = Depends(get_db) +): + """ + Get dashboard statistics for the current user. + Provides overview of documents, chats, and safety metrics. + """ + cache_key = f"dashboard_stats:{user_id}" + + if redis_client: + try: + cached_data = redis_client.get(cache_key) + if cached_data: + return DashboardStats.model_validate_json(cached_data) + except Exception as e: + logger.warning(f"Redis cache read error: {e}") + + now = datetime.utcnow() + today_start = now.replace(hour=0, minute=0, second=0, microsecond=0) + week_start = today_start - timedelta(days=7) + + # Document counts + doc_query = db.query(Document).filter(Document.user_id == user_id) + + total_documents = doc_query.count() + processed_documents = doc_query.filter( + Document.status == DocumentStatus.COMPLETED + ).count() + pending_documents = doc_query.filter( + Document.status.in_( + [ + DocumentStatus.PENDING, + DocumentStatus.PROCESSING, + DocumentStatus.ANALYZING, + ] + ) + ).count() + failed_documents = doc_query.filter( + Document.status == DocumentStatus.FAILED + ).count() + + # Today and week counts + docs_today = doc_query.filter(Document.created_at >= today_start).count() + docs_week = doc_query.filter(Document.created_at >= week_start).count() + + # Chat counts + total_sessions = ( + db.query(ChatSession).filter(ChatSession.user_id == user_id).count() + ) + total_messages = ( + db.query(ChatMessage) + .join(ChatSession) + .filter(ChatSession.user_id == user_id) + .count() + ) + + # Safety metrics + safety_stats = ( + db.query( + func.avg(Document.safety_score).label("avg_score"), + func.count(case((Document.hazards_detected.isnot(None), 1))).label( + "with_hazards" + ), + func.count( + case((Document.compliance_status == ComplianceStatus.VIOLATION, 1)) + ).label("violations"), + func.count( + case((Document.compliance_status == ComplianceStatus.WARNING, 1)) + ).label("warnings"), + ) + .filter( + Document.user_id == user_id, Document.status == DocumentStatus.COMPLETED + ) + .first() + ) + + # Category breakdown + category_counts = ( + db.query(Document.category, func.count(Document.id).label("count")) + .filter(Document.user_id == user_id, Document.category.isnot(None)) + .group_by(Document.category) + .all() + ) + + total_categorized = sum(c.count for c in category_counts) + categories = [ + CategoryCount( + category=cat.value if cat else "other", + count=count, + percentage=( + round(count / total_categorized * 100, 1) + if total_categorized > 0 + else 0 + ), + ) + for cat, count in category_counts + ] + + # Last activity + last_doc = doc_query.order_by(Document.created_at.desc()).first() + last_chat = ( + db.query(ChatSession) + .filter(ChatSession.user_id == user_id) + .order_by(ChatSession.updated_at.desc()) + .first() + ) + + stats = DashboardStats( + total_documents=total_documents, + processed_documents=processed_documents, + pending_documents=pending_documents, + failed_documents=failed_documents, + total_chat_sessions=total_sessions, + total_messages=total_messages, + average_safety_score=( + round(safety_stats.avg_score, 1) if safety_stats.avg_score else None + ), + documents_with_hazards=safety_stats.with_hazards or 0, + compliance_violations=safety_stats.violations or 0, + compliance_warnings=safety_stats.warnings or 0, + documents_processed_today=docs_today, + documents_processed_this_week=docs_week, + documents_by_category=categories, + last_upload_at=last_doc.created_at if last_doc else None, + last_chat_at=last_chat.updated_at if last_chat else None, + ) + + if redis_client: + try: + # Cache for 60 seconds + redis_client.setex(cache_key, 60, stats.model_dump_json()) + except Exception as e: + logger.warning(f"Redis cache write error: {e}") + + return stats + + +@router.get("/documents", response_model=DocumentAnalytics) +async def get_document_analytics( + days: int = Query(30, ge=1, le=365), + user_id: str = Depends(get_current_user_id), + db: Session = Depends(get_db), +): + """ + Get detailed document analytics over time. + """ + start_date = datetime.utcnow() - timedelta(days=days) + + # Uploads by day + uploads = ( + db.query( + func.date(Document.created_at).label("date"), + func.count(Document.id).label("count"), + ) + .filter(Document.user_id == user_id, Document.created_at >= start_date) + .group_by(func.date(Document.created_at)) + .order_by("date") + .all() + ) + + uploads_by_day = [{"date": str(d), "count": c} for d, c in uploads] + + # Category distribution + categories = ( + db.query(Document.category, func.count(Document.id).label("count")) + .filter(Document.user_id == user_id, Document.category.isnot(None)) + .group_by(Document.category) + .all() + ) + + total = sum(c for _, c in categories) + by_category = [ + CategoryCount( + category=cat.value if cat else "other", + count=count, + percentage=round(count / total * 100, 1) if total > 0 else 0, + ) + for cat, count in categories + ] + + # Status distribution + statuses = ( + db.query(Document.status, func.count(Document.id).label("count")) + .filter(Document.user_id == user_id) + .group_by(Document.status) + .all() + ) + + by_status = [ + StatusCount(status=s.value if s else "unknown", count=c) for s, c in statuses + ] + + # File type distribution + file_types = ( + db.query(Document.file_type, func.count(Document.id).label("count")) + .filter(Document.user_id == user_id) + .group_by(Document.file_type) + .all() + ) + + by_file_type = [{"type": ft, "count": c} for ft, c in file_types] + + return DocumentAnalytics( + uploads_by_day=uploads_by_day, + by_category=by_category, + by_status=by_status, + by_file_type=by_file_type, + ) + + +@router.get("/safety", response_model=SafetyAnalytics) +async def get_safety_analytics( + user_id: str = Depends(get_current_user_id), db: Session = Depends(get_db) +): + """ + Get safety compliance analytics. + """ + # Base query for completed documents with safety scores + base = db.query(Document).filter( + Document.user_id == user_id, + Document.status == DocumentStatus.COMPLETED, + Document.safety_score.isnot(None), + ) + + # Score statistics + stats = ( + db.query( + func.avg(Document.safety_score).label("avg"), + func.min(Document.safety_score).label("min"), + func.max(Document.safety_score).label("max"), + ) + .filter(Document.user_id == user_id, Document.safety_score.isnot(None)) + .first() + ) + + # Score distribution + ranges = [ + ("0-25", 0, 25), + ("26-50", 26, 50), + ("51-75", 51, 75), + ("76-100", 76, 100), + ] + + total_with_scores = base.count() + distribution = [] + + for label, low, high in ranges: + count = base.filter( + Document.safety_score >= low, Document.safety_score <= high + ).count() + distribution.append( + SafetyDistribution( + range=label, + count=count, + percentage=( + round(count / total_with_scores * 100, 1) + if total_with_scores > 0 + else 0 + ), + ) + ) + + # Compliance counts + compliant = ( + db.query(Document) + .filter( + Document.user_id == user_id, + Document.compliance_status == ComplianceStatus.COMPLIANT, + ) + .count() + ) + + warnings = ( + db.query(Document) + .filter( + Document.user_id == user_id, + Document.compliance_status == ComplianceStatus.WARNING, + ) + .count() + ) + + violations = ( + db.query(Document) + .filter( + Document.user_id == user_id, + Document.compliance_status == ComplianceStatus.VIOLATION, + ) + .count() + ) + + return SafetyAnalytics( + average_safety_score=round(stats.avg, 1) if stats.avg else 0, + min_safety_score=stats.min, + max_safety_score=stats.max, + score_distribution=distribution, + compliant_count=compliant, + warning_count=warnings, + violation_count=violations, + ) + + +@router.get("/violations") +async def get_recent_violations( + limit: int = Query(20, ge=1, le=100), + user_id: str = Depends(get_current_user_id), + db: Session = Depends(get_db), +): + """ + Get recent violations and warnings from analyzed documents. + Returns documents with detected hazards, ordered by severity and recency. + """ + # Get documents with violations or warnings that have hazards + docs = ( + db.query(Document) + .filter( + Document.user_id == user_id, + Document.status == DocumentStatus.COMPLETED, + Document.compliance_status.in_( + [ComplianceStatus.VIOLATION, ComplianceStatus.WARNING] + ), + Document.hazards_detected.isnot(None), + ) + .order_by(Document.created_at.desc()) + .limit(limit) + .all() + ) + + violations = [] + for doc in docs: + hazards = doc.hazards_detected or [] + if isinstance(hazards, list): + for h in hazards[:3]: # Max 3 hazards per doc + if isinstance(h, dict): + violations.append( + { + "document_id": str(doc.id), + "document_title": doc.title, + "file_name": doc.file_name, + "hazard_type": h.get("type", "Unknown Hazard"), + "severity": h.get("severity", "medium"), + "description": h.get("description", ""), + "regulation": h.get("regulation", ""), + "detected_at": ( + doc.processed_at.isoformat() + if doc.processed_at + else doc.created_at.isoformat() + ), + "compliance_status": ( + doc.compliance_status.value + if doc.compliance_status + else "warning" + ), + } + ) + + return { + "violations": violations, + "total": len(violations), + } diff --git a/app/api/v1/chat.py b/app/api/v1/chat.py new file mode 100644 index 0000000000000000000000000000000000000000..9f98184efb437e9eb1be28bbb4ce51f249788744 --- /dev/null +++ b/app/api/v1/chat.py @@ -0,0 +1,331 @@ +""" +Chat API Endpoints +Chat sessions and AI conversations with RAG +""" + +import logging +import uuid +from datetime import datetime +from typing import List, Optional + +from fastapi import APIRouter, Depends, HTTPException, Query +from sqlalchemy.orm import Session + +from app.api.deps import get_current_user_id +from app.core.exceptions import NotFoundError +from app.db.session import get_db +from app.models.audit import AuditAction, create_audit_log +from app.models.chat import ChatMessage, ChatSession +from app.schemas.chat import ( + ChatMessageResponse, + ChatRequest, + ChatResponse, + ChatSessionCreate, + ChatSessionDetailResponse, + ChatSessionResponse, + ChatSessionUpdateRequest, +) + +logger = logging.getLogger(__name__) + +router = APIRouter() + + +@router.get("/sessions", response_model=List[ChatSessionResponse]) +async def list_chat_sessions( + limit: int = Query(50, ge=1, le=100), + user_id: str = Depends(get_current_user_id), + db: Session = Depends(get_db), +): + """ + List user's chat sessions ordered by most recent. + """ + sessions = ( + db.query(ChatSession) + .filter(ChatSession.user_id == user_id) + .order_by(ChatSession.updated_at.desc()) + .limit(limit) + .all() + ) + + result = [] + for session in sessions: + last_msg = session.messages[-1] if session.messages else None + result.append( + ChatSessionResponse( + id=str(session.id), + title=session.title, + message_count=len(session.messages), + document_context=session.document_context or [], + created_at=session.created_at, + updated_at=session.updated_at, + last_message=last_msg.content[:100] if last_msg else None, + last_message_at=last_msg.created_at if last_msg else None, + ) + ) + + return result + + +@router.post("/sessions", response_model=ChatSessionResponse) +async def create_chat_session( + request: ChatSessionCreate, + user_id: str = Depends(get_current_user_id), + db: Session = Depends(get_db), +): + """ + Create a new chat session. + """ + session = ChatSession( + user_id=user_id, + title=request.title or "New Chat", + document_context=request.document_ids or [], + system_prompt=request.system_prompt, + ) + + db.add(session) + db.commit() + db.refresh(session) + + # Audit log + audit = create_audit_log( + action=AuditAction.CHAT_CREATE.value, + user_id=user_id, + resource_type="chat_session", + resource_id=str(session.id), + ) + db.add(audit) + db.commit() + + return ChatSessionResponse( + id=str(session.id), + title=session.title, + message_count=0, + document_context=session.document_context, + created_at=session.created_at, + updated_at=session.updated_at, + ) + + +@router.get("/sessions/{session_id}", response_model=ChatSessionDetailResponse) +async def get_chat_session( + session_id: uuid.UUID, + user_id: str = Depends(get_current_user_id), + db: Session = Depends(get_db), +): + """ + Get chat session with all messages. + """ + session = ( + db.query(ChatSession) + .filter(ChatSession.id == session_id, ChatSession.user_id == user_id) + .first() + ) + + if not session: + raise NotFoundError("Chat session", session_id) + + messages = [ + ChatMessageResponse( + id=str(msg.id), + role=msg.role, + content=msg.content, + sources=msg.sources or [], + created_at=msg.created_at, + model_used=msg.model_used, + response_time_ms=msg.response_time_ms, + ) + for msg in session.messages + ] + + return ChatSessionDetailResponse( + id=str(session.id), + title=session.title, + document_context=session.document_context or [], + system_prompt=session.system_prompt, + messages=messages, + created_at=session.created_at, + updated_at=session.updated_at, + ) + + +@router.patch("/sessions/{session_id}", response_model=ChatSessionResponse) +async def update_chat_session( + session_id: uuid.UUID, + request: ChatSessionUpdateRequest, + user_id: str = Depends(get_current_user_id), + db: Session = Depends(get_db), +): + """ + Update chat session title or document context. + """ + session = ( + db.query(ChatSession) + .filter(ChatSession.id == session_id, ChatSession.user_id == user_id) + .first() + ) + + if not session: + raise NotFoundError("Chat session", session_id) + + if request.title is not None: + session.title = request.title + if request.document_ids is not None: + session.document_context = request.document_ids + + session.updated_at = datetime.utcnow() + db.commit() + db.refresh(session) + + return ChatSessionResponse( + id=str(session.id), + title=session.title, + message_count=len(session.messages), + document_context=session.document_context, + created_at=session.created_at, + updated_at=session.updated_at, + ) + + +@router.delete("/sessions/{session_id}") +async def delete_chat_session( + session_id: uuid.UUID, + user_id: str = Depends(get_current_user_id), + db: Session = Depends(get_db), +): + """ + Delete a chat session and all messages. + """ + session = ( + db.query(ChatSession) + .filter(ChatSession.id == session_id, ChatSession.user_id == user_id) + .first() + ) + + if not session: + raise NotFoundError("Chat session", session_id) + + # Audit log + audit = create_audit_log( + action=AuditAction.CHAT_DELETE.value, + user_id=user_id, + resource_type="chat_session", + resource_id=session_id, + ) + db.add(audit) + + db.delete(session) + db.commit() + + return {"success": True, "message": "Chat session deleted"} + + +@router.post("/send", response_model=ChatResponse) +async def send_message( + request: ChatRequest, + user_id: str = Depends(get_current_user_id), + db: Session = Depends(get_db), +): + """ + Send a message and get AI response. + Uses RAG to find relevant document context. + """ + start_time = datetime.utcnow() + + # Get or create session + if request.session_id: + session = ( + db.query(ChatSession) + .filter( + ChatSession.id == request.session_id, ChatSession.user_id == user_id + ) + .first() + ) + if not session: + raise NotFoundError("Chat session", request.session_id) + else: + # Create new session + session = ChatSession(user_id=user_id, title="New Chat") + db.add(session) + db.commit() + db.refresh(session) + + # Save user message (but don't commit yet - wait for successful response) + user_message = ChatMessage( + session_id=session.id, role="user", content=request.content + ) + db.add(user_message) + + # Generate AI response with RAG + from app.services.chat_service import ChatService + + chat_service = ChatService() + + try: + ai_response, sources, tokens_used = await chat_service.generate_response( + query=request.content, + user_id=user_id, + document_ids=request.document_ids, + db=db, + ) + + # Calculate response time + end_time = datetime.utcnow() + response_time_ms = int((end_time - start_time).total_seconds() * 1000) + + # Save assistant message + assistant_message = ChatMessage( + session_id=session.id, + role="assistant", + content=ai_response, + sources=sources if request.include_sources else [], + model_used="gemini-2.5-flash", + response_time_ms=response_time_ms, + tokens_used=tokens_used, + ) + db.add(assistant_message) + + # Get fresh message count from database + message_count = ( + db.query(ChatMessage).filter(ChatMessage.session_id == session.id).count() + ) + + # Update session title if first/second message + if message_count <= 2: + # Auto-generate title from first user message + session.title = request.content[:50] + ( + "..." if len(request.content) > 50 else "" + ) + + session.updated_at = datetime.utcnow() + + # Audit log + audit = create_audit_log( + action=AuditAction.CHAT_MESSAGE.value, + user_id=user_id, + resource_type="chat_session", + resource_id=str(session.id), + details={"message_length": len(request.content)}, + ) + db.add(audit) + + # Commit all changes atomically + db.commit() + db.refresh(assistant_message) + except Exception as e: + db.rollback() + raise + + return ChatResponse( + message=ChatMessageResponse( + id=str(assistant_message.id), + role="assistant", + content=ai_response, + sources=sources if request.include_sources else [], + created_at=assistant_message.created_at, + model_used="gemini-2.5-flash", + response_time_ms=response_time_ms, + ), + session_id=str(session.id), + session_title=session.title, + ) diff --git a/app/api/v1/chat_stream.py b/app/api/v1/chat_stream.py new file mode 100644 index 0000000000000000000000000000000000000000..b0584c32082aa255e3fe9034b476fe6f4515b294 --- /dev/null +++ b/app/api/v1/chat_stream.py @@ -0,0 +1,169 @@ +""" +Streaming Chat Endpoint (SSE) +Server-Sent Events streaming for real-time AI responses. + +Delivers three event types to the client: + 1. 'sources' — document citations (emitted first so UI renders immediately) + 2. 'token' — streamed LLM response tokens + 3. 'done' — signals stream completion with metadata + 4. 'error' — on failure +""" + +import json +import logging +from datetime import datetime +from typing import List, Optional + +from fastapi import APIRouter, Depends, Query +from fastapi.responses import StreamingResponse +from sqlalchemy.orm import Session + +from app.api.deps import get_current_user_id +from app.config import settings +from app.db.session import get_db +from app.models.chat import ChatMessage, ChatSession +from app.schemas.chat import ChatRequest + +logger = logging.getLogger(__name__) + +router = APIRouter() + + +@router.post("/stream") +async def stream_chat( + request: ChatRequest, + user_id: str = Depends(get_current_user_id), + db: Session = Depends(get_db), +): + """ + Stream AI response using Server-Sent Events. + + Event format: + event: sources + data: [{"document_title": ..., "file_name": ..., "page_numbers": [...], ...}] + + event: token + data: {"text": "word by word..."} + + event: done + data: {"session_id": "...", "sources_count": 3} + + Usage with fetch(): + const source = new EventSource('/api/v1/chat/stream', {...}) + source.addEventListener('token', (e) => appendToken(JSON.parse(e.data).text)) + source.addEventListener('sources', (e) => renderSources(JSON.parse(e.data))) + source.addEventListener('done', () => source.close()) + """ + from app.services.chat_service import ChatService + + # Get or create session + if request.session_id: + session = ( + db.query(ChatSession) + .filter( + ChatSession.id == request.session_id, ChatSession.user_id == user_id + ) + .first() + ) + if not session: + from app.core.exceptions import NotFoundError + + raise NotFoundError("Chat session", request.session_id) + else: + session = ChatSession(user_id=user_id, title="New Chat") + db.add(session) + db.commit() + db.refresh(session) + + # Save user message immediately + user_message = ChatMessage( + session_id=session.id, + role="user", + content=request.content, + ) + db.add(user_message) + db.commit() + + chat_service = ChatService() + + async def event_generator(): + full_response = [] + sources = [] + tokens_used = None + + try: + async for event in chat_service.generate_response_stream( + query=request.content, + user_id=user_id, + document_ids=request.document_ids, + db=db, + ): + # Forward each SSE event to client + yield event + + # Track sources from the sources event for DB persistence + if event.startswith("event: sources\n"): + data_line = event.split("data: ", 1)[-1].strip() + try: + sources = json.loads(data_line) + except Exception: + pass + + # Accumulate tokens for DB persistence + elif event.startswith("event: token\n"): + data_line = event.split("data: ", 1)[-1].strip() + try: + token_data = json.loads(data_line) + full_response.append(token_data.get("text", "")) + except Exception: + pass + + elif event.startswith("event: done\n"): + data_line = event.split("data: ", 1)[-1].strip() + try: + done_data = json.loads(data_line) + tokens_used = done_data.get("tokens_used") + except Exception: + pass + + except Exception as e: + logger.error(f"Streaming error: {e}", exc_info=True) + yield f"event: error\ndata: {json.dumps({'message': str(e)})}\n\n" + + finally: + # Persist assistant message after stream completes + if full_response: + response_text = "".join(full_response) + assistant_message = ChatMessage( + session_id=session.id, + role="assistant", + content=response_text, + sources=sources if request.include_sources else [], + model_used=settings.GEMINI_MODEL, + tokens_used=tokens_used, + ) + db.add(assistant_message) + + # Auto-title on first message + msg_count = ( + db.query(ChatMessage) + .filter(ChatMessage.session_id == session.id) + .count() + ) + if msg_count <= 2: + session.title = request.content[:50] + ( + "..." if len(request.content) > 50 else "" + ) + + session.updated_at = datetime.utcnow() + db.commit() + + return StreamingResponse( + event_generator(), + media_type="text/event-stream", + headers={ + "Cache-Control": "no-cache", + "Connection": "keep-alive", + "X-Accel-Buffering": "no", # Disable nginx buffering + }, + ) diff --git a/app/api/v1/compliance.py b/app/api/v1/compliance.py new file mode 100644 index 0000000000000000000000000000000000000000..9f9750e3ae008aef50032df017f2f0d4a204ee17 --- /dev/null +++ b/app/api/v1/compliance.py @@ -0,0 +1,265 @@ +""" +Compliance Audit API Endpoints +Regulatory compliance auto-auditor: cross-references operational documents +against regulatory documents to produce per-clause compliance matrices. +""" + +import logging +from typing import Optional +from uuid import UUID + +from fastapi import APIRouter, Depends, HTTPException, Query, status +from sqlalchemy.orm import Session + +from app.api.deps import get_current_user_id +from app.db.session import get_db +from app.models.audit import AuditAction, create_audit_log +from app.models.compliance import AuditStatus, ComplianceAudit, ComplianceMatrixRow +from app.models.document import Document, DocumentCategory, DocumentStatus +from app.schemas.compliance import ( + ComplianceAuditCreate, + ComplianceAuditDetailResponse, + ComplianceAuditListResponse, + ComplianceAuditResponse, + ComplianceMatrixRowResponse, +) +from app.services.queue import enqueue_compliance_task + +logger = logging.getLogger(__name__) + +router = APIRouter() + + +@router.get("/audits", response_model=ComplianceAuditListResponse) +async def list_audits( + page: int = Query(default=1, ge=1), + page_size: int = Query(default=20, ge=1, le=100), + user_id: str = Depends(get_current_user_id), + db: Session = Depends(get_db), +): + """List all compliance audits for the current user.""" + offset = (page - 1) * page_size + + query = ( + db.query(ComplianceAudit) + .filter(ComplianceAudit.user_id == user_id) + .order_by(ComplianceAudit.created_at.desc()) + ) + + total = query.count() + audits = query.offset(offset).limit(page_size).all() + + return ComplianceAuditListResponse( + audits=[ComplianceAuditResponse.model_validate(a) for a in audits], + total=total, + ) + + +@router.post( + "/audits", + response_model=ComplianceAuditResponse, + status_code=status.HTTP_202_ACCEPTED, +) +async def create_audit( + data: ComplianceAuditCreate, + user_id: str = Depends(get_current_user_id), + db: Session = Depends(get_db), +): + """Create and trigger a new compliance audit.""" + # Validate regulation document exists and belongs to user + reg_doc = ( + db.query(Document) + .filter( + Document.id == data.regulation_doc_id, + Document.user_id == user_id, + ) + .first() + ) + if not reg_doc: + raise HTTPException( + status_code=status.HTTP_404_NOT_FOUND, + detail="Regulation document not found", + ) + if reg_doc.status != DocumentStatus.COMPLETED: + raise HTTPException( + status_code=status.HTTP_400_BAD_REQUEST, + detail="Regulation document must be fully processed before auditing", + ) + + # Validate operational documents + op_doc_ids = [str(d) for d in data.operational_doc_ids] + op_docs = ( + db.query(Document) + .filter( + Document.id.in_(op_doc_ids), + Document.user_id == user_id, + ) + .all() + ) + + if len(op_docs) != len(op_doc_ids): + raise HTTPException( + status_code=status.HTTP_400_BAD_REQUEST, + detail="One or more operational documents not found", + ) + + for doc in op_docs: + if doc.status != DocumentStatus.COMPLETED: + raise HTTPException( + status_code=status.HTTP_400_BAD_REQUEST, + detail=f"Document '{doc.title}' must be fully processed first", + ) + + # Create audit record + audit = ComplianceAudit( + user_id=user_id, + title=data.title, + regulation_doc_id=data.regulation_doc_id, + operational_doc_ids=op_doc_ids, + status=AuditStatus.PENDING, + ) + db.add(audit) + db.commit() + db.refresh(audit) + + # Audit log + log = create_audit_log( + user_id=user_id, + action=AuditAction.DOCUMENT_UPLOAD, + resource_type="compliance_audit", + resource_id=str(audit.id), + details={"title": data.title, "regulation_doc": str(data.regulation_doc_id)}, + ) + db.add(log) + db.commit() + + # Enqueue for background processing + await enqueue_compliance_task(str(audit.id)) + + logger.info( + f"Compliance audit created: {audit.id} — " + f"reg_doc={data.regulation_doc_id}, op_docs={len(op_doc_ids)}" + ) + + return ComplianceAuditResponse.model_validate(audit) + + +@router.get("/audits/{audit_id}", response_model=ComplianceAuditDetailResponse) +async def get_audit( + audit_id: UUID, + user_id: str = Depends(get_current_user_id), + db: Session = Depends(get_db), +): + """Get a compliance audit with full matrix rows.""" + audit = ( + db.query(ComplianceAudit) + .filter( + ComplianceAudit.id == audit_id, + ComplianceAudit.user_id == user_id, + ) + .first() + ) + + if not audit: + raise HTTPException( + status_code=status.HTTP_404_NOT_FOUND, + detail="Audit not found", + ) + + rows = ( + db.query(ComplianceMatrixRow) + .filter(ComplianceMatrixRow.audit_id == audit.id) + .order_by(ComplianceMatrixRow.clause_index) + .all() + ) + + return ComplianceAuditDetailResponse( + **ComplianceAuditResponse.model_validate(audit).model_dump(), + rows=[ComplianceMatrixRowResponse.model_validate(r) for r in rows], + ) + + +@router.delete("/audits/{audit_id}") +async def delete_audit( + audit_id: UUID, + user_id: str = Depends(get_current_user_id), + db: Session = Depends(get_db), +): + """Delete a compliance audit and all its matrix rows.""" + audit = ( + db.query(ComplianceAudit) + .filter( + ComplianceAudit.id == audit_id, + ComplianceAudit.user_id == user_id, + ) + .first() + ) + + if not audit: + raise HTTPException( + status_code=status.HTTP_404_NOT_FOUND, + detail="Audit not found", + ) + + audit_title = audit.title + db.delete(audit) + db.commit() + + log = create_audit_log( + user_id=user_id, + action=AuditAction.DOCUMENT_DELETE, + resource_type="compliance_audit", + resource_id=str(audit_id), + details={"title": audit_title}, + ) + db.add(log) + db.commit() + + return {"message": "Audit deleted"} + + +@router.get("/audits/{audit_id}/export") +async def export_audit( + audit_id: UUID, + user_id: str = Depends(get_current_user_id), + db: Session = Depends(get_db), +): + """Export audit results as structured JSON.""" + audit = ( + db.query(ComplianceAudit) + .filter( + ComplianceAudit.id == audit_id, + ComplianceAudit.user_id == user_id, + ) + .first() + ) + + if not audit: + raise HTTPException( + status_code=status.HTTP_404_NOT_FOUND, + detail="Audit not found", + ) + + rows = ( + db.query(ComplianceMatrixRow) + .filter(ComplianceMatrixRow.audit_id == audit.id) + .order_by(ComplianceMatrixRow.clause_index) + .all() + ) + + return { + "audit": audit.to_dict(), + "matrix": [ + { + "clause_index": r.clause_index, + "clause_text": r.clause_text, + "section_title": r.section_title, + "status": r.status, + "assessment": r.assessment, + "confidence": r.confidence, + "evidence_chunks": r.evidence_chunks or [], + "recommendations": r.recommendations or [], + } + for r in rows + ], + } diff --git a/app/api/v1/documents.py b/app/api/v1/documents.py new file mode 100644 index 0000000000000000000000000000000000000000..e7d21427e24fb1429313b8fde16f5fc4e29ceb94 --- /dev/null +++ b/app/api/v1/documents.py @@ -0,0 +1,322 @@ +""" +Document API Endpoints +Document upload, management, and analysis +""" + +import logging +import uuid +from datetime import datetime +from typing import List, Optional + +from fastapi import APIRouter, BackgroundTasks, Depends, HTTPException, Query, status +from sqlalchemy import func +from sqlalchemy.orm import Session + +from app.api.deps import audit_middleware, get_current_user, get_current_user_id +from app.core.exceptions import NotFoundError +from app.db.session import get_db +from app.models.audit import AuditAction, create_audit_log +from app.models.document import ( + ComplianceStatus, + Document, + DocumentCategory, + DocumentStatus, +) +from app.models.user import User +from app.schemas.document import ( + DocumentAnalysisResponse, + DocumentCreate, + DocumentListResponse, + DocumentResponse, + DocumentUploadResponse, +) + +logger = logging.getLogger(__name__) + +router = APIRouter() + + +@router.get("", response_model=DocumentListResponse) +async def list_documents( + page: int = Query(1, ge=1), + page_size: int = Query(20, ge=1, le=100), + category: Optional[DocumentCategory] = None, + status: Optional[DocumentStatus] = None, + search: Optional[str] = None, + user_id: str = Depends(get_current_user_id), + db: Session = Depends(get_db), +): + """ + List user's documents with filtering and pagination. + """ + query = db.query(Document).filter(Document.user_id == user_id) + + # Apply filters + if category: + query = query.filter(Document.category == category) + if status: + query = query.filter(Document.status == status) + if search: + # Escape SQL wildcard characters to prevent injection + escaped_search = ( + search.replace("\\", "\\\\").replace("%", "\\%").replace("_", "\\_") + ) + safe_pattern = f"%{escaped_search}%" + query = query.filter( + Document.title.ilike(safe_pattern, escape="\\") + | Document.file_name.ilike(safe_pattern, escape="\\") + ) + + # Get total count + total = query.count() + + # Apply pagination + offset = (page - 1) * page_size + documents = ( + query.order_by(Document.created_at.desc()).offset(offset).limit(page_size).all() + ) + + # Calculate stats + stats = _calculate_document_stats(db, user_id) + + return DocumentListResponse( + documents=[DocumentResponse(**doc.to_dict()) for doc in documents], + total=total, + page=page, + page_size=page_size, + stats=stats, + ) + + +@router.post( + "", response_model=DocumentUploadResponse, status_code=status.HTTP_202_ACCEPTED +) +async def create_document( + request: DocumentCreate, + user_id: str = Depends(get_current_user_id), + db: Session = Depends(get_db), +): + """ + Create a new document from UploadThing URL. + Triggers background processing with AI analysis via task queue. + """ + # Create document record + document = Document( + user_id=user_id, + title=request.title or request.file_name.rsplit(".", 1)[0], + file_name=request.file_name, + file_size=request.file_size, + file_type=request.file_type, + file_url=request.file_url, + status=DocumentStatus.PENDING, + tags=request.tags or [], + ) + + db.add(document) + db.flush() # Flush to get the document.id without committing + + # Create audit log + audit = create_audit_log( + action=AuditAction.DOCUMENT_UPLOAD.value, + user_id=user_id, + resource_type="document", + resource_id=str(document.id), + details={ + "file_name": document.file_name, + "file_size": document.file_size, + "file_type": document.file_type, + }, + ) + db.add(audit) + + # Commit both document and audit atomically + db.commit() + db.refresh(document) + + # Trigger background processing + from app.services.queue import enqueue_document_task + + enqueue_document_task(str(document.id)) + + logger.info(f"Document created and enqueued: {document.id} - {document.title}") + + return DocumentUploadResponse( + id=str(document.id), + title=document.title, + file_name=document.file_name, + status=DocumentStatus.PENDING, + job_id=str(document.id), # Using doc ID as job ID for now + message="Document uploaded successfully. AI analysis queued.", + ) + + +@router.get("/{document_id}", response_model=DocumentResponse) +async def get_document( + document_id: uuid.UUID, + user_id: str = Depends(get_current_user_id), + db: Session = Depends(get_db), +): + """ + Get document details by ID. + """ + document = ( + db.query(Document) + .filter(Document.id == document_id, Document.user_id == user_id) + .first() + ) + + if not document: + raise NotFoundError("Document", document_id) + + return DocumentResponse(**document.to_dict()) + + +@router.delete("/{document_id}") +async def delete_document( + document_id: uuid.UUID, + user_id: str = Depends(get_current_user_id), + db: Session = Depends(get_db), +): + """ + Delete a document and all associated data. + """ + document = ( + db.query(Document) + .filter(Document.id == document_id, Document.user_id == user_id) + .first() + ) + + if not document: + raise NotFoundError("Document", document_id) + + # Create audit log before deletion + audit = create_audit_log( + action=AuditAction.DOCUMENT_DELETE.value, + user_id=user_id, + resource_type="document", + resource_id=str(document_id), + details={"file_name": document.file_name}, + ) + db.add(audit) + + # Delete document (cascade will handle embeddings) + db.delete(document) + db.commit() + + logger.info(f"Document deleted: {document_id}") + + return {"success": True, "message": "Document deleted successfully"} + + +@router.get("/{document_id}/analysis", response_model=DocumentAnalysisResponse) +async def get_document_analysis( + document_id: uuid.UUID, + user_id: str = Depends(get_current_user_id), + db: Session = Depends(get_db), +): + """ + Get AI analysis results for a document. + """ + document = ( + db.query(Document) + .filter(Document.id == document_id, Document.user_id == user_id) + .first() + ) + + if not document: + raise NotFoundError("Document", document_id) + + if document.status != DocumentStatus.COMPLETED: + return DocumentAnalysisResponse( + document_id=str(document.id), status=document.status.value, analysis=None + ) + + analysis = { + "category": document.category, + "subcategory": document.subcategory, + "classification_confidence": document.classification_confidence, + "summary": document.summary, + "key_points": document.key_points or [], + "safety_score": document.safety_score, + "compliance_status": document.compliance_status, + "hazards_detected": document.hazards_detected or [], + "safety_recommendations": document.safety_recommendations or [], + "entities": { + k: v if isinstance(v, list) else [] + for k, v in (document.entities or {}).items() + }, + } + + return DocumentAnalysisResponse( + document_id=str(document.id), status="completed", analysis=analysis + ) + + +@router.post("/{document_id}/reanalyze", status_code=status.HTTP_202_ACCEPTED) +async def reanalyze_document( + document_id: uuid.UUID, + user_id: str = Depends(get_current_user_id), + db: Session = Depends(get_db), +): + """ + Trigger re-analysis of a document. + """ + document = ( + db.query(Document) + .filter(Document.id == document_id, Document.user_id == user_id) + .first() + ) + + if not document: + raise NotFoundError("Document", document_id) + + # Reset status + document.status = DocumentStatus.PENDING + db.commit() + + # Trigger reprocessing + from app.services.queue import enqueue_document_task + + enqueue_document_task(str(document.id)) + + return {"success": True, "message": "Document reanalysis queued"} + + +def _calculate_document_stats(db: Session, user_id: str) -> dict: + """Calculate aggregated statistics for user's documents""" + # Category distribution + category_counts = ( + db.query(Document.category, func.count(Document.id)) + .filter(Document.user_id == user_id, Document.category.isnot(None)) + .group_by(Document.category) + .all() + ) + + by_category = { + cat.value if cat else "other": count for cat, count in category_counts + } + + # Status distribution + status_counts = ( + db.query(Document.status, func.count(Document.id)) + .filter(Document.user_id == user_id) + .group_by(Document.status) + .all() + ) + + by_status = { + status.value if status else "unknown": count for status, count in status_counts + } + + # Average safety score + avg_score = ( + db.query(func.avg(Document.safety_score)) + .filter(Document.user_id == user_id, Document.safety_score.isnot(None)) + .scalar() + ) + + return { + "by_category": by_category, + "by_status": by_status, + "avg_safety_score": round(avg_score, 2) if avg_score else None, + } diff --git a/app/api/v1/health.py b/app/api/v1/health.py new file mode 100644 index 0000000000000000000000000000000000000000..65a83bec972e1114a144fd14dfa58a0936b5ea08 --- /dev/null +++ b/app/api/v1/health.py @@ -0,0 +1,80 @@ +""" +Health Check Endpoints +System health and status monitoring +""" + +from datetime import datetime + +from fastapi import APIRouter, Depends +from sqlalchemy.orm import Session + +from app.config import settings +from app.db.session import check_db_connection, get_db +from app.schemas.common import HealthResponse + +router = APIRouter() + + +@router.get("/", response_model=HealthResponse) +async def root(): + """Root endpoint - basic health check""" + return HealthResponse( + status="healthy", + version=settings.APP_VERSION, + environment=settings.ENVIRONMENT, + timestamp=datetime.utcnow(), + ) + + +@router.get("/health", response_model=HealthResponse) +async def health_check(db: Session = Depends(get_db)): + """ + Detailed health check with service status. + Checks database connectivity and other services. + """ + services = {"database": "unknown", "redis": "unknown", "ai": "unknown"} + + # Check database + try: + from sqlalchemy import text + + db.execute(text("SELECT 1")) + services["database"] = "healthy" + except Exception as e: + services["database"] = f"unhealthy: {str(e)}" + + # Check Redis (if configured) + r = None + try: + import redis + + r = redis.from_url(settings.REDIS_URL) + r.ping() + services["redis"] = "healthy" + except Exception: + services["redis"] = "not_configured" + finally: + if r is not None: + r.close() + + # Check AI service (Gemini) + try: + import google.generativeai as genai + + genai.configure(api_key=settings.GEMINI_API_KEY) + services["ai"] = "healthy" + except Exception: + services["ai"] = "not_configured" + + # Overall status + overall = "healthy" + if services["database"] != "healthy": + overall = "degraded" + + return HealthResponse( + status=overall, + version=settings.APP_VERSION, + environment=settings.ENVIRONMENT, + timestamp=datetime.utcnow(), + services=services, + ) diff --git a/app/api/v1/jobs.py b/app/api/v1/jobs.py new file mode 100644 index 0000000000000000000000000000000000000000..9a3cfff87545591f202af0d10922575fd5192e7c --- /dev/null +++ b/app/api/v1/jobs.py @@ -0,0 +1,122 @@ +""" +Jobs API Endpoints +Background job status tracking +""" + +import logging + +from fastapi import APIRouter, Depends, Query +from sqlalchemy.orm import Session + +from app.api.deps import get_current_user_id +from app.core.exceptions import NotFoundError +from app.db.session import get_db +from app.models.document import Document, DocumentStatus +from app.schemas.common import JobStatusResponse + +logger = logging.getLogger(__name__) + +router = APIRouter() + + +@router.get("/{job_id}", response_model=JobStatusResponse) +async def get_job_status( + job_id: str, + user_id: str = Depends(get_current_user_id), + db: Session = Depends(get_db), +): + """ + Get status of a background processing job. + Currently jobs are tracked via document ID. + """ + # For now, job_id is document_id + document = ( + db.query(Document) + .filter(Document.id == job_id, Document.user_id == user_id) + .first() + ) + + if not document: + raise NotFoundError("Job", job_id) + + # Map document status to job status + status_map = { + DocumentStatus.PENDING: "pending", + DocumentStatus.PROCESSING: "processing", + DocumentStatus.ANALYZING: "processing", + DocumentStatus.COMPLETED: "completed", + DocumentStatus.FAILED: "failed", + } + + # Calculate progress based on status + progress_map = { + DocumentStatus.PENDING: 0, + DocumentStatus.PROCESSING: 30, + DocumentStatus.ANALYZING: 70, + DocumentStatus.COMPLETED: 100, + DocumentStatus.FAILED: 0, + } + + result = None + if document.status == DocumentStatus.COMPLETED: + result = { + "document_id": str(document.id), + "category": document.category.value if document.category else None, + "safety_score": document.safety_score, + "summary_preview": document.summary[:200] if document.summary else None, + } + + return JobStatusResponse( + job_id=str(document.id), + status=status_map.get(document.status, "unknown"), + progress=progress_map.get(document.status, 0), + result=result, + error=document.processing_error, + created_at=document.created_at, + updated_at=document.updated_at, + completed_at=document.processed_at, + ) + + +@router.get("") +async def list_active_jobs( + user_id: str = Depends(get_current_user_id), db: Session = Depends(get_db) +): + """ + List all active (non-completed) processing jobs. + """ + active_docs = ( + db.query(Document) + .filter( + Document.user_id == user_id, + Document.status.in_( + [ + DocumentStatus.PENDING, + DocumentStatus.PROCESSING, + DocumentStatus.ANALYZING, + ] + ), + ) + .order_by(Document.created_at.desc()) + .all() + ) + + jobs = [] + for doc in active_docs: + progress = { + DocumentStatus.PENDING: 0, + DocumentStatus.PROCESSING: 30, + DocumentStatus.ANALYZING: 70, + }.get(doc.status, 0) + + jobs.append( + { + "job_id": str(doc.id), + "document_title": doc.title, + "status": doc.status.value, + "progress": progress, + "created_at": doc.created_at.isoformat(), + } + ) + + return {"jobs": jobs, "count": len(jobs)} diff --git a/app/api/v1/prompts.py b/app/api/v1/prompts.py new file mode 100644 index 0000000000000000000000000000000000000000..0aa91b1b737fe6b897947328fc1e671c94182ebe --- /dev/null +++ b/app/api/v1/prompts.py @@ -0,0 +1,200 @@ +""" +Custom Prompts API Endpoints +CRUD operations for user-defined AI prompts +""" + +import logging +from typing import List, Optional + +from fastapi import APIRouter, Depends, Query +from sqlalchemy.orm import Session + +from app.api.deps import get_current_user, get_current_user_id +from app.core.exceptions import NotFoundError +from app.db.session import get_db +from app.models.prompt import CustomPrompt +from app.schemas.prompt import ( + PromptCreate, + PromptListResponse, + PromptResponse, + PromptUpdate, +) + +logger = logging.getLogger(__name__) + +router = APIRouter() + + +@router.get("", response_model=List[PromptResponse]) +async def list_prompts( + category: Optional[str] = None, + user_id: str = Depends(get_current_user_id), + db: Session = Depends(get_db), +): + """ + List all custom prompts for the current user. + """ + query = db.query(CustomPrompt).filter(CustomPrompt.user_id == user_id) + + if category: + query = query.filter(CustomPrompt.category == category) + + prompts = query.order_by(CustomPrompt.created_at.desc()).all() + + return [ + PromptResponse( + id=str(p.id), + name=p.name, + prompt=p.prompt_text, + description=p.description, + category=p.category, + is_default=p.is_default, + created_at=p.created_at, + updated_at=p.updated_at, + ) + for p in prompts + ] + + +@router.post("", response_model=PromptResponse, status_code=201) +async def create_prompt( + request: PromptCreate, + user=Depends(get_current_user), + db: Session = Depends(get_db), +): + """ + Create a new custom prompt. + """ + prompt = CustomPrompt( + user_id=user.clerk_user_id, + name=request.name, + prompt_text=request.prompt, + description=request.description, + category=request.category, + is_default=False, + ) + + db.add(prompt) + db.commit() + db.refresh(prompt) + + logger.info(f"Prompt created: {prompt.id} by user {user.clerk_user_id[:12]}") + + return PromptResponse( + id=str(prompt.id), + name=prompt.name, + prompt=prompt.prompt_text, + description=prompt.description, + category=prompt.category, + is_default=prompt.is_default, + created_at=prompt.created_at, + updated_at=prompt.updated_at, + ) + + +@router.get("/{prompt_id}", response_model=PromptResponse) +async def get_prompt( + prompt_id: str, + user_id: str = Depends(get_current_user_id), + db: Session = Depends(get_db), +): + """ + Get a specific prompt by ID. + """ + prompt = ( + db.query(CustomPrompt) + .filter( + CustomPrompt.id == prompt_id, + CustomPrompt.user_id == user_id, + ) + .first() + ) + + if not prompt: + raise NotFoundError("Prompt", prompt_id) + + return PromptResponse( + id=str(prompt.id), + name=prompt.name, + prompt=prompt.prompt_text, + description=prompt.description, + category=prompt.category, + is_default=prompt.is_default, + created_at=prompt.created_at, + updated_at=prompt.updated_at, + ) + + +@router.put("/{prompt_id}", response_model=PromptResponse) +async def update_prompt( + prompt_id: str, + request: PromptUpdate, + user_id: str = Depends(get_current_user_id), + db: Session = Depends(get_db), +): + """ + Update a custom prompt. + """ + prompt = ( + db.query(CustomPrompt) + .filter( + CustomPrompt.id == prompt_id, + CustomPrompt.user_id == user_id, + ) + .first() + ) + + if not prompt: + raise NotFoundError("Prompt", prompt_id) + + if request.name is not None: + prompt.name = request.name + if request.prompt is not None: + prompt.prompt_text = request.prompt + if request.description is not None: + prompt.description = request.description + if request.category is not None: + prompt.category = request.category + + db.commit() + db.refresh(prompt) + + return PromptResponse( + id=str(prompt.id), + name=prompt.name, + prompt=prompt.prompt_text, + description=prompt.description, + category=prompt.category, + is_default=prompt.is_default, + created_at=prompt.created_at, + updated_at=prompt.updated_at, + ) + + +@router.delete("/{prompt_id}") +async def delete_prompt( + prompt_id: str, + user_id: str = Depends(get_current_user_id), + db: Session = Depends(get_db), +): + """ + Delete a custom prompt. + """ + prompt = ( + db.query(CustomPrompt) + .filter( + CustomPrompt.id == prompt_id, + CustomPrompt.user_id == user_id, + ) + .first() + ) + + if not prompt: + raise NotFoundError("Prompt", prompt_id) + + db.delete(prompt) + db.commit() + + logger.info(f"Prompt deleted: {prompt_id}") + + return {"success": True, "message": "Prompt deleted successfully"} diff --git a/app/api/v1/router.py b/app/api/v1/router.py new file mode 100644 index 0000000000000000000000000000000000000000..bfd136d912247b0ae6a3ce3ad516bbc041764f7d --- /dev/null +++ b/app/api/v1/router.py @@ -0,0 +1,42 @@ +""" +API v1 Router +Combines all endpoint routers into single API router +""" + +from fastapi import APIRouter + +from app.api.v1 import ( + analytics, + chat, + chat_stream, + compliance, + documents, + health, + jobs, + prompts, + search, + user, +) + +api_router = APIRouter() + +# Include all routers +api_router.include_router(health.router, tags=["Health"]) + +api_router.include_router(documents.router, prefix="/documents", tags=["Documents"]) + +api_router.include_router(chat.router, prefix="/chat", tags=["Chat"]) + +api_router.include_router(chat_stream.router, prefix="/chat", tags=["Chat"]) + +api_router.include_router(analytics.router, prefix="/analytics", tags=["Analytics"]) + +api_router.include_router(jobs.router, prefix="/jobs", tags=["Jobs"]) + +api_router.include_router(prompts.router, prefix="/prompts", tags=["Prompts"]) + +api_router.include_router(user.router, prefix="/user", tags=["User"]) + +api_router.include_router(search.router, prefix="/search", tags=["Search"]) + +api_router.include_router(compliance.router, prefix="/compliance", tags=["Compliance"]) diff --git a/app/api/v1/search.py b/app/api/v1/search.py new file mode 100644 index 0000000000000000000000000000000000000000..579d33aed415116107f8b8b95f16d53ffe13ca8b --- /dev/null +++ b/app/api/v1/search.py @@ -0,0 +1,138 @@ +""" +Semantic Search API Endpoint +Natural language search across all user documents. + +Production retrieval pipeline: + 1. Embed query via Gemini text-embedding-004 + 2. Hybrid search: pgvector cosine + pg_trgm BM25 via RRF + 3. Cross-encoder reranking for precise relevance + 4. Results grouped by document with page-level provenance +""" + +import asyncio +import logging +from typing import List, Optional + +import google.generativeai as genai +from fastapi import APIRouter, Depends +from fastapi import Query as FastAPIQuery +from sqlalchemy.orm import Session + +from app.api.deps import get_current_user_id +from app.config import settings +from app.db.session import get_db + +logger = logging.getLogger(__name__) + +router = APIRouter() + +# Configure Gemini for embeddings +genai.configure(api_key=settings.GEMINI_API_KEY) + + +async def _get_query_embedding(text_input: str) -> List[float]: + """Generate query embedding using Gemini text-embedding-004.""" + try: + result = await asyncio.to_thread( + genai.embed_content, + model=settings.EMBEDDING_MODEL, + content=text_input, + task_type="retrieval_query", + ) + return result["embedding"] + except Exception as e: + logger.error(f"Query embedding failed: {e}") + return [] + + +@router.get("") +async def semantic_search( + q: str = FastAPIQuery( + ..., min_length=2, description="Natural language search query" + ), + limit: int = FastAPIQuery(default=20, ge=1, le=50), + category: Optional[str] = FastAPIQuery(default=None), + user_id: str = Depends(get_current_user_id), + db: Session = Depends(get_db), +): + """ + Semantic search across all user documents. + + Production pipeline: + 1. Hybrid search (pgvector + pg_trgm via RRF) + 2. Cross-encoder reranking + 3. Results with page-level provenance + + Example queries: + - "ventilation rules for underground mines" + - "30 CFR 75.323 methane requirements" + - "equipment maintenance schedule for Caterpillar D11" + """ + if not q.strip(): + return {"query": q, "results": [], "total": 0} + + # Generate query embedding + embedding = await _get_query_embedding(q.strip()) + if not embedding: + return { + "query": q, + "results": [], + "total": 0, + "error": "Could not generate embedding for query", + } + + # Step 1: Hybrid search (over-fetch for reranking) + from app.services.hybrid_search import hybrid_search + + candidates = await hybrid_search( + query_text=q.strip(), + query_embedding=embedding, + db=db, + user_id=user_id, + top_k=settings.RERANK_OVER_FETCH, + ) + + # Step 2: Rerank + if settings.ENABLE_RERANKING and len(candidates) > limit: + from app.services.reranker import rerank + + candidates = rerank(query=q.strip(), chunks=candidates, top_k=limit) + else: + candidates = candidates[:limit] + + # Format results + results = [] + for row in candidates: + pages = row.get("page_numbers", []) + page_str = ( + f"Pages {pages[0]}\u2013{pages[-1]}" + if len(pages) > 1 + else f"Page {pages[0]}" if pages else "Unknown page" + ) + results.append( + { + "chunk_id": row["id"], + "document_id": row["document_id"], + "document_title": row["document_title"], + "file_name": row["file_name"], + "chunk_text": row["text"], + "section_title": row.get("section_title"), + "page_numbers": pages, + "page_label": page_str, + "relevance_score": round( + row.get("rerank_score", row.get("score", 0.0)), 4 + ), + "relevance_percent": round( + row.get("rerank_score", row.get("score", 0.0)) * 100, 1 + ), + } + ) + + return { + "query": q, + "results": results, + "total": len(results), + "filters_applied": { + "category": category, + }, + } diff --git a/app/api/v1/user.py b/app/api/v1/user.py new file mode 100644 index 0000000000000000000000000000000000000000..f6e30144916304458ec617229869787eb94ca153 --- /dev/null +++ b/app/api/v1/user.py @@ -0,0 +1,105 @@ +""" +User Profile API Endpoints +User profile management endpoints +""" + +import logging +from typing import List, Optional + +from fastapi import APIRouter, Depends +from pydantic import BaseModel +from sqlalchemy.orm import Session + +from app.api.deps import get_current_user, get_current_user_id +from app.db.session import get_db +from app.models.user import User + +logger = logging.getLogger(__name__) + +router = APIRouter() + + +class UserProfileResponse(BaseModel): + """User profile response""" + + clerk_user_id: str + email: Optional[str] = None + full_name: Optional[str] = None + avatar_url: Optional[str] = None + company_name: Optional[str] = None + company_role: Optional[str] = None + industry_focus: Optional[List[str]] = None + mine_sites: Optional[List[str]] = None + is_active: bool = True + + class Config: + from_attributes = True + + +class UserProfileUpdate(BaseModel): + """User profile update request""" + + full_name: Optional[str] = None + company_name: Optional[str] = None + company_role: Optional[str] = None + industry_focus: Optional[List[str]] = None + mine_sites: Optional[List[str]] = None + + +@router.get("/profile", response_model=UserProfileResponse) +async def get_user_profile( + user: User = Depends(get_current_user), +): + """ + Get the current user's profile. + Creates the user record on first access if it doesn't exist. + """ + return UserProfileResponse( + clerk_user_id=user.clerk_user_id, + email=user.email, + full_name=user.full_name, + avatar_url=user.avatar_url, + company_name=user.company_name, + company_role=user.company_role, + industry_focus=user.industry_focus, + mine_sites=user.mine_sites, + is_active=user.is_active, + ) + + +@router.put("/profile", response_model=UserProfileResponse) +async def update_user_profile( + request: UserProfileUpdate, + user: User = Depends(get_current_user), + db: Session = Depends(get_db), +): + """ + Update the current user's profile. + """ + if request.full_name is not None: + user.full_name = request.full_name + if request.company_name is not None: + user.company_name = request.company_name + if request.company_role is not None: + user.company_role = request.company_role + if request.industry_focus is not None: + user.industry_focus = request.industry_focus + if request.mine_sites is not None: + user.mine_sites = request.mine_sites + + db.commit() + db.refresh(user) + + logger.info(f"User profile updated: {user.clerk_user_id[:12]}") + + return UserProfileResponse( + clerk_user_id=user.clerk_user_id, + email=user.email, + full_name=user.full_name, + avatar_url=user.avatar_url, + company_name=user.company_name, + company_role=user.company_role, + industry_focus=user.industry_focus, + mine_sites=user.mine_sites, + is_active=user.is_active, + ) diff --git a/app/config.py b/app/config.py new file mode 100644 index 0000000000000000000000000000000000000000..7b576c01bf12eed2ad883b79146e5eca072547f2 --- /dev/null +++ b/app/config.py @@ -0,0 +1,125 @@ +""" +Application Configuration +Centralized settings management using Pydantic Settings +""" + +import os +from functools import lru_cache +from typing import List, Optional + +from pydantic import Field +from pydantic_settings import BaseSettings, SettingsConfigDict + + +class Settings(BaseSettings): + """Application settings loaded from environment variables""" + + # Application + APP_NAME: str = "MiningNiti" + APP_VERSION: str = "2.0.0" + DEBUG: bool = Field(default=False) + ENVIRONMENT: str = Field(default="development") + + # API + API_V1_PREFIX: str = "/api/v1" + CORS_ORIGINS: List[str] = Field( + default=["http://localhost:3000", "https://*.vercel.app"] + ) + + # Database + DATABASE_URL: str = Field(..., description="PostgreSQL connection string") + DB_POOL_SIZE: int = Field(default=5) + DB_MAX_OVERFLOW: int = Field(default=10) + + # Redis + REDIS_URL: str = Field(default="redis://localhost:6379/0") + + # AI/ML - Multi-Provider Setup + GEMINI_API_KEY: str = Field(..., description="Google Gemini API Key") + GROQ_API_KEY: str = Field( + ..., description="Groq API Key for Classifier & Entity Extractors" + ) + MISTRAL_API_KEY: str = Field(..., description="Mistral API Key for Safety Analyzer") + CEREBRAS_API_KEY: str = Field(default="", description="Cerebras API Key") + + GEMINI_MODEL: str = Field(default="gemini-1.5-flash") + EMBEDDING_MODEL: str = Field(default="models/gemini-embedding-001") + + AGENT_PROVIDER_MAP: dict = { + "embeddings": {"provider": "gemini", "model": "text-embedding-004"}, + "chat_service": {"provider": "gemini", "model": "gemini-1.5-flash"}, + "summarizer_agent": {"provider": "gemini", "model": "gemini-1.5-flash"}, + "classifier_agent": {"provider": "groq", "model": "llama-3.3-70b-versatile"}, + "entity_extractor": {"provider": "cerebras", "model": "llama-4-scout"}, + "safety_analyzer": {"provider": "mistral", "model": "magistral-small-latest"}, + "fallback": {"provider": "openrouter", "model": "deepseek/deepseek-r1:free"}, + } + + # Authentication - Clerk + CLERK_JWKS_URL: str = Field(..., description="Clerk JWKS URL for JWT verification") + + # Document Processing + MAX_FILE_SIZE_MB: int = Field(default=50) + ALLOWED_FILE_TYPES: List[str] = Field( + default=[ + "application/pdf", + "application/vnd.openxmlformats-officedocument.wordprocessingml.document", + "text/plain", + ] + ) + CHUNK_SIZE: int = Field(default=1000) + CHUNK_OVERLAP: int = Field(default=200) + + # Mining AI Settings + SAFETY_SCORE_THRESHOLD: float = Field(default=70.0) + MAX_EMBEDDINGS_PER_QUERY: int = Field(default=5) + + # RAG Pipeline — Production Retrieval + RERANK_MODEL: str = Field( + default="cross-encoder/ms-marco-MiniLM-L-6-v2", + description="Cross-encoder model for reranking retrieved chunks", + ) + RERANK_OVER_FETCH: int = Field( + default=20, + description="How many chunks to fetch from vector+BM25 before reranking", + ) + RERANK_TOP_K: int = Field( + default=5, + description="Final number of chunks after reranking", + ) + SIMILARITY_THRESHOLD: float = Field( + default=0.25, + description="Minimum cosine similarity to include a chunk (0-1)", + ) + ENABLE_HYBRID_SEARCH: bool = Field( + default=True, + description="Combine vector search with pg_trgm BM25 via RRF", + ) + ENABLE_RERANKING: bool = Field( + default=True, + description="Apply cross-encoder reranking after retrieval", + ) + RRF_K: int = Field( + default=60, + description="Reciprocal Rank Fusion constant (higher = less rank influence)", + ) + + # SSL + SSL_CERT_PATH: Optional[str] = Field(default=None) + + model_config = SettingsConfigDict( + env_file=".env", env_file_encoding="utf-8", case_sensitive=True, extra="ignore" + ) + + +@lru_cache() +def get_settings() -> Settings: + """ + Get cached settings instance. + Uses lru_cache for performance - settings are loaded once. + """ + return Settings() + + +# Convenience export +settings = get_settings() diff --git a/app/core/__init__.py b/app/core/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..5ed1bca5cb577ca7cb5dd7b7e2c2c2bd0062ea5a --- /dev/null +++ b/app/core/__init__.py @@ -0,0 +1,3 @@ +""" +Core module - Security, Exceptions, Utilities +""" diff --git a/app/core/exceptions.py b/app/core/exceptions.py new file mode 100644 index 0000000000000000000000000000000000000000..f9a1fd9b6ae8bd8504ffd9530e42967b35106dad --- /dev/null +++ b/app/core/exceptions.py @@ -0,0 +1,115 @@ +""" +Custom Exception Classes +Enterprise-grade error handling with proper HTTP status codes +""" + +from typing import Any, Dict, Optional + +from fastapi import HTTPException, status + + +class MiningNitiException(Exception): + """Base exception for all MiningNiti errors""" + + def __init__( + self, + message: str, + code: str = "INTERNAL_ERROR", + details: Optional[Dict[str, Any]] = None, + ): + self.message = message + self.code = code + self.details = details or {} + super().__init__(self.message) + + +class AuthenticationError(HTTPException): + """Raised when authentication fails""" + + def __init__(self, detail: str = "Authentication failed"): + import logging + + logging.getLogger("app.core.exceptions").warning( + f"AuthenticationError raised: {detail}" + ) + super().__init__( + status_code=status.HTTP_401_UNAUTHORIZED, + detail=detail, + headers={"WWW-Authenticate": "Bearer"}, + ) + + +class AuthorizationError(HTTPException): + """Raised when user lacks permission""" + + def __init__(self, detail: str = "Permission denied"): + super().__init__(status_code=status.HTTP_403_FORBIDDEN, detail=detail) + + +class NotFoundError(HTTPException): + """Raised when resource is not found""" + + def __init__(self, resource: str = "Resource", resource_id: str = ""): + detail = f"{resource} not found" + if resource_id: + detail = f"{resource} with id '{resource_id}' not found" + super().__init__(status_code=status.HTTP_404_NOT_FOUND, detail=detail) + + +class ValidationError(HTTPException): + """Raised when request validation fails""" + + def __init__( + self, detail: str = "Validation failed", errors: Optional[list] = None + ): + super().__init__( + status_code=status.HTTP_422_UNPROCESSABLE_ENTITY, + detail={"message": detail, "errors": errors or []}, + ) + + +class DocumentProcessingError(MiningNitiException): + """Raised when document processing fails""" + + def __init__(self, message: str, document_id: Optional[str] = None): + super().__init__( + message=message, + code="DOCUMENT_PROCESSING_ERROR", + details={"document_id": document_id} if document_id else {}, + ) + + +class AIServiceError(MiningNitiException): + """Raised when AI service (Gemini) fails""" + + def __init__(self, message: str, service: str = "gemini"): + super().__init__( + message=message, code="AI_SERVICE_ERROR", details={"service": service} + ) + + +class RateLimitError(HTTPException): + """Raised when rate limit is exceeded""" + + def __init__(self, retry_after: int = 60): + super().__init__( + status_code=status.HTTP_429_TOO_MANY_REQUESTS, + detail=f"Rate limit exceeded. Retry after {retry_after} seconds.", + headers={"Retry-After": str(retry_after)}, + ) + + +class JobNotFoundError(NotFoundError): + """Raised when background job is not found""" + + def __init__(self, job_id: str): + super().__init__(resource="Job", resource_id=job_id) + + +class SafetyViolationError(MiningNitiException): + """Raised when safety compliance check fails critically""" + + def __init__(self, message: str, violations: list): + super().__init__( + message=message, code="SAFETY_VIOLATION", details={"violations": violations} + ) diff --git a/app/core/security.py b/app/core/security.py new file mode 100644 index 0000000000000000000000000000000000000000..4020b427203bd981ffe45a95982091e0bf4255fb --- /dev/null +++ b/app/core/security.py @@ -0,0 +1,148 @@ +""" +Security Module +JWT verification, authentication, and authorization utilities +""" + +import logging +from datetime import datetime, timedelta +from functools import lru_cache +from typing import Any, Dict, Optional + +import httpx +from jose import JWTError, jwk, jwt +from jose.exceptions import ExpiredSignatureError + +from app.config import settings +from app.core.exceptions import AuthenticationError + +logger = logging.getLogger(__name__) + + +class JWKSClient: + """ + JWKS (JSON Web Key Set) client for Clerk JWT verification. + Caches keys to avoid repeated network requests. + """ + + def __init__(self, jwks_url: str): + self.jwks_url = jwks_url + self._keys: Dict[str, Any] = {} + self._last_fetch: Optional[datetime] = None + self._cache_duration = timedelta(hours=1) + + async def get_signing_key(self, kid: str) -> Optional[Dict[str, Any]]: + """Get signing key by key ID (kid)""" + await self._refresh_keys_if_needed() + return self._keys.get(kid) + + async def _refresh_keys_if_needed(self, force: bool = False): + """Refresh keys if cache is stale or force is True""" + now = datetime.utcnow() + + if ( + not force + and self._last_fetch + and (now - self._last_fetch) < self._cache_duration + ): + return + + try: + async with httpx.AsyncClient() as client: + response = await client.get(self.jwks_url, timeout=10.0) + response.raise_for_status() + jwks_data = response.json() + + self._keys = {key["kid"]: key for key in jwks_data.get("keys", [])} + self._last_fetch = now + logger.info(f"Refreshed JWKS keys: {len(self._keys)} keys loaded") + + except Exception as e: + logger.error(f"Failed to fetch JWKS: {e}") + if not self._keys: + raise AuthenticationError("Unable to verify authentication") + + +# Global JWKS client instance +_jwks_client: Optional[JWKSClient] = None + + +def get_jwks_client() -> JWKSClient: + """Get or create JWKS client singleton""" + global _jwks_client + if _jwks_client is None: + _jwks_client = JWKSClient(settings.CLERK_JWKS_URL) + return _jwks_client + + +async def verify_jwt_token(token: str) -> Dict[str, Any]: + """ + Verify JWT token from Clerk. + + Args: + token: JWT token string + + Returns: + Decoded token payload with user information + + Raises: + AuthenticationError: If token is invalid or expired + """ + try: + # Decode header to get key ID + unverified_header = jwt.get_unverified_header(token) + kid = unverified_header.get("kid") + + if not kid: + raise AuthenticationError("Invalid token: missing key ID") + + # Get signing key + jwks_client = get_jwks_client() + signing_key = await jwks_client.get_signing_key(kid) + + if not signing_key: + # Force refresh and try again + await jwks_client._refresh_keys_if_needed(force=True) + signing_key = await jwks_client.get_signing_key(kid) + + if not signing_key: + raise AuthenticationError("Invalid token: unknown signing key") + + # Verify and decode token + payload = jwt.decode( + token, + signing_key, + algorithms=["RS256"], + options={"verify_aud": False}, # Clerk doesn't always set audience + ) + + return payload + + except ExpiredSignatureError: + raise AuthenticationError("Token has expired") + except JWTError as e: + logger.warning(f"JWT verification failed: {e}") + raise AuthenticationError("Invalid token") + except Exception as e: + logger.error(f"Authentication error: {e}") + raise AuthenticationError("Authentication failed") + + +def extract_user_id(payload: Dict[str, Any]) -> str: + """ + Extract user ID from JWT payload. + Clerk uses 'sub' claim for user ID. + """ + user_id = payload.get("sub") + if not user_id: + raise AuthenticationError("Invalid token: missing user ID") + return user_id + + +def extract_user_email(payload: Dict[str, Any]) -> Optional[str]: + """Extract email from JWT payload if available""" + # Clerk may include email in different claims + return ( + payload.get("email") + or payload.get("primary_email") + or payload.get("email_addresses", [{}])[0].get("email_address") + ) diff --git a/app/db/__init__.py b/app/db/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..67fc7790a0182d691f6032b93cc72ac8ca1a5e13 --- /dev/null +++ b/app/db/__init__.py @@ -0,0 +1,7 @@ +""" +Database module - Session management and migrations +""" + +from app.db.session import Base, SessionLocal, engine, get_db + +__all__ = ["engine", "SessionLocal", "get_db", "Base"] diff --git a/app/db/session.py b/app/db/session.py new file mode 100644 index 0000000000000000000000000000000000000000..4eea25872c02b9d0cbe4d889b5a2472f03fbcadd --- /dev/null +++ b/app/db/session.py @@ -0,0 +1,117 @@ +""" +Database Session Management +SQLAlchemy engine and session configuration with connection pooling +""" + +import logging +from contextlib import contextmanager +from typing import Generator + +from sqlalchemy import create_engine, event, text +from sqlalchemy.orm import Session, sessionmaker +from sqlalchemy.pool import QueuePool + +from app.config import settings + +# Use the single canonical Base so all models share the same metadata +from app.models.base import Base # noqa: F401 - re-exported for convenience + +logger = logging.getLogger(__name__) + +# Configure engine with connection pooling +engine_args = { + "pool_size": settings.DB_POOL_SIZE, + "max_overflow": settings.DB_MAX_OVERFLOW, + "pool_pre_ping": True, # Verify connections before use + "pool_recycle": 3600, # Recycle connections after 1 hour + "echo": False, # Suppress excessive SQL query logging +} + +# Add SSL config if certificate path provided +connect_args = {} +if settings.SSL_CERT_PATH: + connect_args["sslmode"] = "require" + connect_args["sslrootcert"] = settings.SSL_CERT_PATH + +engine = create_engine( + settings.DATABASE_URL, poolclass=QueuePool, connect_args=connect_args, **engine_args +) + +# Session factory +SessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=engine) + + +# Connection event listeners for debugging +@event.listens_for(engine, "connect") +def on_connect(dbapi_conn, connection_record): + logger.debug("Database connection established") + + +@event.listens_for(engine, "checkout") +def on_checkout(dbapi_conn, connection_record, connection_proxy): + logger.debug("Database connection checked out from pool") + + +def get_db() -> Generator[Session, None, None]: + """ + Dependency for FastAPI endpoints. + Yields a database session and ensures cleanup. + + Usage: + @app.get("/items") + def get_items(db: Session = Depends(get_db)): + ... + """ + db = SessionLocal() + try: + yield db + finally: + db.close() + + +@contextmanager +def get_db_context() -> Generator[Session, None, None]: + """ + Context manager for database sessions outside of FastAPI. + Useful for background workers and scripts. + + Usage: + with get_db_context() as db: + db.query(...) + """ + db = SessionLocal() + try: + yield db + db.commit() + except Exception: + db.rollback() + raise + finally: + db.close() + + +def init_db(): + """Initialize database tables using the canonical Base from models.base""" + # Import all models so they register their tables with Base.metadata + from app.models import audit, chat, document, prompt, user # noqa: F401 + from app.models.base import Base as ModelBase + + # Ensure pgvector extension exists before creating tables that use VECTOR columns + with engine.connect() as conn: + conn.execute(text("CREATE EXTENSION IF NOT EXISTS vector")) + conn.commit() + logger.info("pgvector extension ensured") + + ModelBase.metadata.create_all(bind=engine) + logger.info("Database tables created successfully") + + +def check_db_connection() -> bool: + """Check if database is reachable""" + try: + with engine.connect() as conn: + conn.execute(text("SELECT 1")) + return True + except Exception as e: + logger.error(f"Database connection failed: {e}") + return False diff --git a/app/main.py b/app/main.py new file mode 100644 index 0000000000000000000000000000000000000000..e5d77657e915f5e632d4131fde9b497b98d35b9e --- /dev/null +++ b/app/main.py @@ -0,0 +1,283 @@ +""" +MiningNiti Enterprise Backend +FastAPI Application Entry Point + +AI-Powered Document Intelligence for the Coal Mining Industry +""" + +import logging +from contextlib import asynccontextmanager +from datetime import datetime + +from fastapi import FastAPI, HTTPException, Request, status +from fastapi.exceptions import RequestValidationError +from fastapi.middleware.cors import CORSMiddleware +from fastapi.responses import JSONResponse +from slowapi import Limiter, _rate_limit_exceeded_handler +from slowapi.errors import RateLimitExceeded +from slowapi.middleware import SlowAPIMiddleware +from slowapi.util import get_remote_address + +from app.api.v1 import api_router +from app.config import settings +from app.core.exceptions import MiningNitiException +from app.db.session import check_db_connection, init_db + +# Configure logging +logging.basicConfig( + level=logging.DEBUG if settings.DEBUG else logging.INFO, + format="%(asctime)s - %(name)s - %(levelname)s - %(message)s", +) +logger = logging.getLogger(__name__) + +# Suppress SQLAlchemy's extremely verbose SQL echo in debug mode — +# it drowns out real application logs. Set to WARNING to only see errors. +logging.getLogger("sqlalchemy.engine").setLevel(logging.WARNING) +logging.getLogger("sqlalchemy.pool").setLevel(logging.WARNING) +logging.getLogger("sqlalchemy.dialects").setLevel(logging.WARNING) +# Also suppress httpcore connection-level debug spam +logging.getLogger("httpcore").setLevel(logging.WARNING) +logging.getLogger("httpx").setLevel(logging.INFO) + + +@asynccontextmanager +async def lifespan(app: FastAPI): + """Application lifespan handler for startup and shutdown events""" + # Startup + logger.info(f"Starting {settings.APP_NAME} v{settings.APP_VERSION}") + logger.info(f"Environment: {settings.ENVIRONMENT}") + + import asyncio + + from app.services.queue import compliance_worker, document_worker + + worker_task = asyncio.create_task(document_worker()) + compliance_worker_task = asyncio.create_task(compliance_worker()) + + # Check database connection + if check_db_connection(): + logger.info("Database connection verified") + # Auto-create tables on startup (idempotent) + try: + init_db() + logger.info("Database tables initialized") + except Exception as e: + logger.warning(f"Database table creation warning: {e}") + + # Recovery: reset documents stuck in transient states from a previous crash/restart + try: + from app.db.session import get_db_context + from app.models.document import Document, DocumentStatus + + with get_db_context() as db: + stuck_docs = ( + db.query(Document) + .filter(Document.status.in_(["processing", "analyzing"])) + .all() + ) + if stuck_docs: + for doc in stuck_docs: + doc.status = DocumentStatus.PENDING + doc.processing_error = "Reset after server restart" + db.commit() + logger.info( + f"Recovery: reset {len(stuck_docs)} stuck document(s) to PENDING" + ) + else: + logger.info("Recovery: no stuck documents found") + except Exception as e: + logger.warning(f"Document recovery warning: {e}") + + # Recovery: reset compliance audits stuck in running state + try: + from app.models.compliance import AuditStatus, ComplianceAudit + + with get_db_context() as db: + stuck_audits = ( + db.query(ComplianceAudit) + .filter(ComplianceAudit.status.in_(["running"])) + .all() + ) + if stuck_audits: + for audit in stuck_audits: + audit.status = AuditStatus.PENDING + audit.processing_error = "Reset after server restart" + db.commit() + logger.info( + f"Recovery: reset {len(stuck_audits)} stuck audit(s) to PENDING" + ) + except Exception as e: + logger.warning(f"Audit recovery warning: {e}") + else: + logger.warning("Database connection failed - some features may not work") + + yield + + # Shutdown + logger.info("Shutting down application") + worker_task.cancel() + compliance_worker_task.cancel() + try: + await worker_task + except asyncio.CancelledError: + pass + try: + await compliance_worker_task + except asyncio.CancelledError: + pass + + +# Create FastAPI application +app = FastAPI( + title=settings.APP_NAME, + description=""" +## MiningNiti - AI Document Intelligence for Mining + +Enterprise-grade document processing and AI chat platform +specifically designed for the coal mining industry. + +### Features +- 📄 **Smart Document Processing** - Upload PDF, DOCX, TXT with AI analysis +- 🤖 **Multi-Agent AI** - Classification, Safety Analysis, Entity Extraction +- 💬 **RAG Chat** - Context-aware conversations with document citations +- 📊 **Analytics Dashboard** - Safety metrics, compliance tracking +- 🔒 **Enterprise Security** - JWT auth, audit logging + +### AI Agents +1. **Classifier Agent** - Categorizes mining documents +2. **Safety Analyzer** - Detects hazards and compliance issues +3. **Entity Extractor** - Extracts equipment, chemicals, regulations +4. **Summarizer** - Creates executive summaries + """, + version=settings.APP_VERSION, + docs_url="/docs", + redoc_url="/redoc", + openapi_url="/openapi.json", + lifespan=lifespan, +) + +# Rate Limiter +limiter = Limiter(key_func=get_remote_address, default_limits=["120/minute"]) + +# CORS Configuration +_EXTRA_ORIGINS = ["http://localhost:3000", "http://localhost:3001"] + +app.state.limiter = limiter +app.add_exception_handler(RateLimitExceeded, _rate_limit_exceeded_handler) +app.add_middleware(SlowAPIMiddleware) + +app.add_middleware( + CORSMiddleware, + allow_origins=settings.CORS_ORIGINS + _EXTRA_ORIGINS, + allow_credentials=True, + allow_methods=["*"], + allow_headers=["*"], +) + + +# Exception Handlers + + +def _get_cors_headers(request: Request) -> dict: + """ + Build CORS headers to attach to error responses. + This is needed because FastAPI's HTTPBearer can short-circuit before + CORSMiddleware has a chance to add Access-Control-Allow-Origin headers, + causing the browser to report a CORS error instead of the real auth error. + """ + origin = request.headers.get("origin", "") + allowed_origins = settings.CORS_ORIGINS + _EXTRA_ORIGINS + if origin in allowed_origins or any( + origin.endswith(o.lstrip("*")) for o in allowed_origins if "*" in o + ): + return { + "Access-Control-Allow-Origin": origin, + "Access-Control-Allow-Credentials": "true", + } + return {} + + +@app.exception_handler(HTTPException) +async def http_exception_handler(request: Request, exc: HTTPException): + """Handle HTTP exceptions with CORS headers so auth failures are visible to the browser""" + headers = {**(exc.headers or {}), **_get_cors_headers(request)} + return JSONResponse( + status_code=exc.status_code, + content={"detail": exc.detail}, + headers=headers, + ) + + +@app.exception_handler(MiningNitiException) +async def miningniti_exception_handler(request: Request, exc: MiningNitiException): + """Handle custom application exceptions""" + return JSONResponse( + status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, + content={ + "error": exc.message, + "code": exc.code, + "details": exc.details, + "timestamp": datetime.utcnow().isoformat(), + }, + headers=_get_cors_headers(request), + ) + + +@app.exception_handler(RequestValidationError) +async def validation_exception_handler(request: Request, exc: RequestValidationError): + """Handle Pydantic validation errors""" + return JSONResponse( + status_code=status.HTTP_422_UNPROCESSABLE_ENTITY, + content={ + "error": "Validation failed", + "code": "VALIDATION_ERROR", + "details": exc.errors(), + "timestamp": datetime.utcnow().isoformat(), + }, + headers=_get_cors_headers(request), + ) + + +@app.exception_handler(Exception) +async def global_exception_handler(request: Request, exc: Exception): + """Global exception handler for unhandled errors""" + logger.error(f"Unhandled exception: {exc}", exc_info=True) + return JSONResponse( + status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, + content={ + "error": "Internal server error", + "code": "INTERNAL_SERVER_ERROR", + "timestamp": datetime.utcnow().isoformat(), + }, + headers=_get_cors_headers(request), + ) + + +# Include API router +app.include_router(api_router, prefix=settings.API_V1_PREFIX) + + +# Root endpoint (without /api/v1 prefix for health checks) +@app.get("/", tags=["Root"]) +async def root(): + """Root endpoint - application info""" + return { + "name": settings.APP_NAME, + "version": settings.APP_VERSION, + "description": "AI Document Intelligence for Mining Industry", + "docs": "/docs", + "health": f"{settings.API_V1_PREFIX}/health", + } + + +@app.get("/health", tags=["Root"]) +async def health(): + """Quick health check for load balancers""" + return {"status": "healthy", "timestamp": datetime.utcnow().isoformat()} + + +# Run with: uvicorn app.main:app --reload +if __name__ == "__main__": + import uvicorn + + uvicorn.run("app.main:app", host="0.0.0.0", port=8000, reload=settings.DEBUG) diff --git a/app/models/__init__.py b/app/models/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..08b00dfb383f0c3d87f6590191fcbd548f51491e --- /dev/null +++ b/app/models/__init__.py @@ -0,0 +1,30 @@ +""" +Database Models +SQLAlchemy ORM models for the MiningNiti platform +""" + +from app.models.audit import AuditAction, AuditLog +from app.models.base import Base, TimestampMixin, UUIDMixin +from app.models.chat import ChatMessage, ChatSession +from app.models.compliance import AuditStatus, ComplianceAudit, ComplianceMatrixRow +from app.models.document import Document, DocumentCategory, DocumentEmbedding +from app.models.prompt import CustomPrompt +from app.models.user import User + +__all__ = [ + "Base", + "TimestampMixin", + "UUIDMixin", + "User", + "Document", + "DocumentEmbedding", + "DocumentCategory", + "ChatSession", + "ChatMessage", + "AuditLog", + "AuditAction", + "CustomPrompt", + "ComplianceAudit", + "ComplianceMatrixRow", + "AuditStatus", +] diff --git a/app/models/audit.py b/app/models/audit.py new file mode 100644 index 0000000000000000000000000000000000000000..204854356327dc345005c1461ecdec26c74caf8f --- /dev/null +++ b/app/models/audit.py @@ -0,0 +1,142 @@ +""" +Audit Log Model +Enterprise compliance logging for all user actions +""" + +import uuid +from datetime import datetime +from enum import Enum + +from sqlalchemy import JSON, Column, DateTime, String, Text +from sqlalchemy.dialects.postgresql import JSONB as PG_JSONB +from sqlalchemy.dialects.postgresql import UUID + +JSONB = JSON().with_variant(PG_JSONB, "postgresql") + +from app.models.base import Base, UUIDMixin + + +class AuditAction(str, Enum): + """Types of auditable actions""" + + # Document actions + DOCUMENT_UPLOAD = "document.upload" + DOCUMENT_VIEW = "document.view" + DOCUMENT_DELETE = "document.delete" + DOCUMENT_PROCESS = "document.process" + + # Chat actions + CHAT_CREATE = "chat.create" + CHAT_MESSAGE = "chat.message" + CHAT_DELETE = "chat.delete" + + # User actions + USER_LOGIN = "user.login" + USER_LOGOUT = "user.logout" + USER_PROFILE_UPDATE = "user.profile_update" + + # Admin actions + ADMIN_ACTION = "admin.action" + + # System actions + SYSTEM_ERROR = "system.error" + AI_ANALYSIS = "ai.analysis" + + +class AuditLog(Base, UUIDMixin): + """ + Immutable audit log for enterprise compliance. + Tracks all user actions and system events. + """ + + __tablename__ = "audit_logs" + + # Who + user_id = Column( + String(255), nullable=True, index=True + ) # Nullable for system events + user_email = Column(String(255), nullable=True) + + # What + action = Column(String(100), nullable=False, index=True) + resource_type = Column(String(50), nullable=True) # document, chat, user + resource_id = Column(String(255), nullable=True) + + # Details + description = Column(Text, nullable=True) + details = Column(JSONB, default={}) + # Example details: + # { + # "file_name": "safety_manual.pdf", + # "file_size": 1024000, + # "category": "safety_protocol" + # } + + # When + timestamp = Column(DateTime, default=datetime.utcnow, nullable=False, index=True) + + # Where + ip_address = Column(String(50), nullable=True) + user_agent = Column(Text, nullable=True) + + # Outcome + success = Column(String(10), default="true") # true, false, partial + error_message = Column(Text, nullable=True) + + def __repr__(self): + return f"" + + def to_dict(self): + """Convert to dictionary for API responses""" + return { + "id": str(self.id), + "user_id": self.user_id, + "action": self.action, + "resource_type": self.resource_type, + "resource_id": self.resource_id, + "description": self.description, + "details": self.details, + "timestamp": self.timestamp.isoformat() if self.timestamp else None, + "success": self.success, + } + + +def create_audit_log( + action: str, + user_id: str = None, + user_email: str = None, + resource_type: str = None, + resource_id: str = None, + description: str = None, + details: dict = None, + ip_address: str = None, + user_agent: str = None, + success: str = "true", + error_message: str = None, +) -> AuditLog: + """ + Factory function to create audit log entries. + + Usage: + log = create_audit_log( + action=AuditAction.DOCUMENT_UPLOAD.value, + user_id=current_user.id, + resource_type="document", + resource_id=str(doc.id), + details={"file_name": doc.file_name} + ) + db.add(log) + """ + return AuditLog( + action=action, + user_id=user_id, + user_email=user_email, + resource_type=resource_type, + resource_id=resource_id, + description=description, + details=details or {}, + ip_address=ip_address, + user_agent=user_agent, + success=success, + error_message=error_message, + ) diff --git a/app/models/base.py b/app/models/base.py new file mode 100644 index 0000000000000000000000000000000000000000..6e3572283d0deda204199491c9fac0a03bfb9a14 --- /dev/null +++ b/app/models/base.py @@ -0,0 +1,41 @@ +""" +Base Model Classes +Mixins and base classes for all SQLAlchemy models +""" + +import uuid +from datetime import datetime + +from sqlalchemy import Column, DateTime +from sqlalchemy.dialects.postgresql import UUID +from sqlalchemy.ext.declarative import declarative_base, declared_attr + +Base = declarative_base() + + +class UUIDMixin: + """Mixin that adds a UUID primary key""" + + id = Column(UUID(as_uuid=True), primary_key=True, default=uuid.uuid4, index=True) + + +class TimestampMixin: + """Mixin that adds created_at and updated_at timestamps""" + + created_at = Column(DateTime, default=datetime.utcnow, nullable=False) + + updated_at = Column( + DateTime, default=datetime.utcnow, onupdate=datetime.utcnow, nullable=True + ) + + +class TableNameMixin: + """Mixin that auto-generates table name from class name""" + + @declared_attr + def __tablename__(cls): + # Convert CamelCase to snake_case + name = cls.__name__ + return "".join(["_" + c.lower() if c.isupper() else c for c in name]).lstrip( + "_" + ) diff --git a/app/models/chat.py b/app/models/chat.py new file mode 100644 index 0000000000000000000000000000000000000000..b5ce675a80a1b81218ff7655b3020d9b4568ac71 --- /dev/null +++ b/app/models/chat.py @@ -0,0 +1,126 @@ +""" +Chat Models +Chat sessions and messages with RAG context +""" + +import uuid +from datetime import datetime + +from sqlalchemy import JSON, Column, DateTime, ForeignKey, Integer, String, Text +from sqlalchemy.dialects.postgresql import JSONB as PG_JSONB +from sqlalchemy.dialects.postgresql import UUID + +JSONB = JSON().with_variant(PG_JSONB, "postgresql") +from sqlalchemy.orm import relationship + +from app.models.base import Base, TimestampMixin, UUIDMixin + + +class ChatSession(Base, UUIDMixin, TimestampMixin): + """ + Chat session for grouping related messages. + Each session maintains context for the conversation. + """ + + __tablename__ = "chat_sessions" + + # Owner + user_id = Column( + String(255), ForeignKey("users.clerk_user_id"), nullable=False, index=True + ) + + # Session info + title = Column(String(500), nullable=False, default="New Chat") + + # Context - selected documents for this session + document_context = Column(JSONB, default=list) # List of document IDs + + # Custom prompt if set + system_prompt = Column(Text, nullable=True) + + # Metadata + metadata_ = Column("metadata", JSONB, default=dict) + + # Relationships + user = relationship("User", back_populates="chat_sessions") + messages = relationship( + "ChatMessage", + back_populates="session", + cascade="all, delete-orphan", + order_by="ChatMessage.created_at", + ) + + @property + def message_count(self) -> int: + return len(self.messages) + + def __repr__(self): + return f"" + + def to_dict(self): + """Convert to dictionary for API responses""" + return { + "id": str(self.id), + "title": self.title, + "message_count": self.message_count, + "document_context": self.document_context, + "created_at": self.created_at.isoformat() if self.created_at else None, + "updated_at": self.updated_at.isoformat() if self.updated_at else None, + } + + +class ChatMessage(Base, UUIDMixin): + """ + Individual chat message with RAG source citations. + """ + + __tablename__ = "chat_messages" + + # Parent session + session_id = Column( + UUID(as_uuid=True), + ForeignKey("chat_sessions.id", ondelete="CASCADE"), + nullable=False, + index=True, + ) + + # Message content + role = Column(String(20), nullable=False) # "user" or "assistant" + content = Column(Text, nullable=False) + + # RAG sources - documents/chunks used for this response + sources = Column(JSONB, default=[]) + # Structure: + # [ + # { + # "document_id": "uuid", + # "document_title": "Safety Manual", + # "chunk_text": "...", + # "relevance_score": 0.95, + # "page": 5 + # } + # ] + + # AI metadata + model_used = Column(String(100), nullable=True) + tokens_used = Column(JSONB, nullable=True) # {"input": 100, "output": 50} + response_time_ms = Column(Integer, nullable=True) + + # Timestamp + created_at = Column(DateTime, default=datetime.utcnow, nullable=False) + + # Relationships + session = relationship("ChatSession", back_populates="messages") + + def __repr__(self): + return f"" + + def to_dict(self): + """Convert to dictionary for API responses""" + return { + "id": str(self.id), + "role": self.role, + "content": self.content, + "sources": self.sources, + "created_at": self.created_at.isoformat() if self.created_at else None, + } diff --git a/app/models/compliance.py b/app/models/compliance.py new file mode 100644 index 0000000000000000000000000000000000000000..054c9474fe86efca891afa8a9c10e767da8b7de2 --- /dev/null +++ b/app/models/compliance.py @@ -0,0 +1,156 @@ +""" +Compliance Audit Models +Regulatory compliance auto-auditor: cross-references operational documents +against regulatory documents (MSHA/OSHA/EPA/DGMS) to produce per-clause +compliance matrices with citations. +""" + +import uuid +from datetime import datetime, timezone +from enum import Enum + +from sqlalchemy import JSON, Column, DateTime +from sqlalchemy import Enum as SQLEnum +from sqlalchemy import Float, ForeignKey, Integer, String, Text +from sqlalchemy.dialects.postgresql import JSONB as PG_JSONB +from sqlalchemy.dialects.postgresql import UUID + +JSONB = JSON().with_variant(PG_JSONB, "postgresql") +from sqlalchemy.orm import relationship + +from app.models.base import Base, TimestampMixin, UUIDMixin + + +class AuditStatus(str, Enum): + """Compliance audit processing status""" + + PENDING = "pending" + RUNNING = "running" + COMPLETED = "completed" + FAILED = "failed" + + +class ComplianceAudit(Base, UUIDMixin, TimestampMixin): + """ + A compliance audit that cross-references one regulatory document + against a set of operational documents. + """ + + __tablename__ = "compliance_audits" + + # Owner + user_id = Column( + String(255), + ForeignKey("users.clerk_user_id"), + nullable=False, + index=True, + ) + + title = Column(String(500), nullable=False) + + # The regulatory document being audited against + regulation_doc_id = Column( + UUID(as_uuid=True), + ForeignKey("documents.id"), + nullable=False, + ) + + # List of operational document UUIDs being audited + operational_doc_ids = Column(JSONB, nullable=False, default=list) + + # Status + status = Column( + SQLEnum(AuditStatus), + default=AuditStatus.PENDING, + nullable=False, + index=True, + ) + + # Aggregate stats + total_clauses = Column(Integer, nullable=True) + processed_clauses = Column(Integer, default=0, nullable=False) + compliant_count = Column(Integer, nullable=True) + gap_count = Column(Integer, nullable=True) + missing_count = Column(Integer, nullable=True) + overall_score = Column(Float, nullable=True) # 0-100 + + processing_error = Column(Text, nullable=True) + completed_at = Column(DateTime, nullable=True) + + # Relationships + rows = relationship( + "ComplianceMatrixRow", + back_populates="audit", + cascade="all, delete-orphan", + ) + regulation_doc = relationship( + "Document", + foreign_keys=[regulation_doc_id], + ) + + def __repr__(self): + return f"" + + def to_dict(self): + return { + "id": str(self.id), + "title": self.title, + "regulation_doc_id": str(self.regulation_doc_id), + "operational_doc_ids": [str(d) for d in (self.operational_doc_ids or [])], + "status": self.status.value if self.status else None, + "total_clauses": self.total_clauses, + "processed_clauses": self.processed_clauses, + "compliant_count": self.compliant_count, + "gap_count": self.gap_count, + "missing_count": self.missing_count, + "overall_score": self.overall_score, + "processing_error": self.processing_error, + "completed_at": ( + self.completed_at.replace(tzinfo=timezone.utc).isoformat() + if self.completed_at + else None + ), + "created_at": ( + self.created_at.replace(tzinfo=timezone.utc).isoformat() + if self.created_at + else None + ), + } + + +class ComplianceMatrixRow(Base, UUIDMixin): + """ + A single row in the compliance matrix: one regulation clause + assessed against operational document evidence. + """ + + __tablename__ = "compliance_matrix_rows" + + audit_id = Column( + UUID(as_uuid=True), + ForeignKey("compliance_audits.id", ondelete="CASCADE"), + nullable=False, + index=True, + ) + + clause_index = Column(Integer, nullable=False) + clause_text = Column(Text, nullable=False) + section_title = Column(String(500), nullable=True) + + # compliant | gap | missing + status = Column(String(50), nullable=False) + assessment = Column(Text, nullable=False) # LLM explanation + confidence = Column(Float, nullable=False) # 0.0-1.0 + + # Evidence chunks that informed the assessment + evidence_chunks = Column(JSONB, nullable=True) + # Structure: [{"chunk_text": "...", "document_title": "...", + # "page_numbers": [12,13], "relevance_score": 0.87}] + + recommendations = Column(JSONB, nullable=True) # list of strings + + # Relationship + audit = relationship("ComplianceAudit", back_populates="rows") + + def __repr__(self): + return f"" diff --git a/app/models/document.py b/app/models/document.py new file mode 100644 index 0000000000000000000000000000000000000000..37b1115c88f2a2f06f674e28d0ec562341888614 --- /dev/null +++ b/app/models/document.py @@ -0,0 +1,230 @@ +""" +Document Models +Document storage, classification, and embeddings +""" + +import uuid +from datetime import datetime, timezone +from enum import Enum + +from sqlalchemy import JSON, Column, DateTime +from sqlalchemy import Enum as SQLEnum +from sqlalchemy import Float, ForeignKey, Integer, String, Text +from sqlalchemy.dialects.postgresql import ARRAY +from sqlalchemy.dialects.postgresql import JSONB as PG_JSONB +from sqlalchemy.dialects.postgresql import UUID + +JSONB = JSON().with_variant(PG_JSONB, "postgresql") +from pgvector.sqlalchemy import Vector +from sqlalchemy.orm import relationship + +from app.models.base import Base, TimestampMixin, UUIDMixin + + +class DocumentCategory(str, Enum): + """Mining document categories for classification""" + + SAFETY_PROTOCOL = "safety_protocol" + EQUIPMENT_MANUAL = "equipment_manual" + REGULATORY = "regulatory" + INCIDENT_REPORT = "incident_report" + GEOLOGICAL = "geological" + ENVIRONMENTAL = "environmental" + TRAINING = "training" + PERMIT = "permit" + MAINTENANCE = "maintenance" + OTHER = "other" + + +class DocumentStatus(str, Enum): + """Document processing status""" + + PENDING = "pending" + PROCESSING = "processing" + ANALYZING = "analyzing" + COMPLETED = "completed" + FAILED = "failed" + + +class ComplianceStatus(str, Enum): + """Safety compliance status""" + + COMPLIANT = "compliant" + WARNING = "warning" + VIOLATION = "violation" + PENDING = "pending" + NOT_APPLICABLE = "not_applicable" + + +class Document(Base, UUIDMixin, TimestampMixin): + """ + Document model with AI-enhanced metadata. + Stores file info, classification, safety analysis, and extracted entities. + """ + + __tablename__ = "documents" + + # Owner + user_id = Column( + String(255), ForeignKey("users.clerk_user_id"), nullable=False, index=True + ) + + # File information + title = Column(String(500), nullable=False) + file_name = Column(String(500), nullable=False) + file_size = Column(Integer, nullable=False) # bytes + file_type = Column(String(100), nullable=False) # MIME type + file_url = Column(Text, nullable=False) + + # Processing status + status = Column( + SQLEnum(DocumentStatus), + default=DocumentStatus.PENDING, + nullable=False, + index=True, + ) + processing_error = Column(Text, nullable=True) + processed_at = Column(DateTime, nullable=True) + + # Extracted content + content = Column(Text, nullable=True) # Full text content + page_count = Column(Integer, nullable=True) # deprecated alias — use total_pages + total_pages = Column( + Integer, nullable=True + ) # authoritative page count from extractor + word_count = Column(Integer, nullable=True) + + # AI Classification + category = Column( + SQLEnum(DocumentCategory), + default=DocumentCategory.OTHER, + nullable=True, + index=True, + ) + subcategory = Column(String(100), nullable=True) + classification_confidence = Column(Float, nullable=True) # 0.0 - 1.0 + + # AI Summary + summary = Column(Text, nullable=True) # AI-generated summary + key_points = Column(JSONB, nullable=True) # List of key points + + # Safety Analysis + safety_score = Column(Float, nullable=True) # 0-100 + compliance_status = Column( + SQLEnum(ComplianceStatus), default=ComplianceStatus.PENDING, nullable=True + ) + hazards_detected = Column(JSONB, nullable=True) # List of hazards + safety_recommendations = Column(JSONB, nullable=True) + + # Named Entity Recognition + entities = Column(JSONB, nullable=True) + # Structure: + # { + # "equipment": ["Caterpillar D11", "Komatsu PC8000"], + # "chemicals": ["methane", "coal dust"], + # "locations": ["Mine Site A", "Section 4B"], + # "personnel": ["John Smith", "Safety Team"], + # "dates": ["2024-01-15", "Q1 2024"], + # "regulations": ["MSHA 30 CFR 75.400", "OSHA 1910.134"] + # } + + # Extra Metadata + extra_metadata = Column("metadata", JSONB, default=dict) + tags = Column(JSONB, default=list) + + # Relationships + user = relationship("User", back_populates="documents") + embeddings = relationship( + "DocumentEmbedding", back_populates="document", cascade="all, delete-orphan" + ) + + def __repr__(self): + return f"" + + def to_dict(self): + """Convert to dictionary for API responses""" + return { + "id": str(self.id), + "title": self.title, + "file_name": self.file_name, + "file_size": self.file_size, + "file_type": self.file_type, + "file_url": self.file_url, + "status": self.status.value if self.status else None, + "category": self.category.value if self.category else None, + "subcategory": self.subcategory, + "classification_confidence": self.classification_confidence, + "summary": self.summary, + "key_points": self.key_points, + "safety_score": self.safety_score, + "compliance_status": ( + self.compliance_status.value if self.compliance_status else None + ), + "hazards_detected": self.hazards_detected, + "entities": { + k: v if isinstance(v, list) else [] + for k, v in (self.entities or {}).items() + } + or None, + "page_count": self.page_count, + "word_count": self.word_count, + "created_at": ( + self.created_at.replace(tzinfo=timezone.utc).isoformat() + if self.created_at + else None + ), + "processed_at": ( + self.processed_at.replace(tzinfo=timezone.utc).isoformat() + if self.processed_at + else None + ), + "total_pages": self.total_pages or self.page_count, + } + + +class DocumentEmbedding(Base, UUIDMixin): + """ + Vector embeddings for document chunks. + Used for semantic search and RAG. + + The embedding column uses pgvector's native Vector(768) type with an + HNSW index (see migration 001) for sub-5ms approximate nearest-neighbor + search instead of brute-force Python cosine similarity. + """ + + __tablename__ = "document_embeddings" + + # Parent document + document_id = Column( + UUID(as_uuid=True), + ForeignKey("documents.id", ondelete="CASCADE"), + nullable=False, + index=True, + ) + + # Chunk information + chunk_index = Column(Integer, nullable=False) + chunk_text = Column(Text, nullable=False) + + # Vector embedding — native pgvector type with HNSW index (see migration 001) + # Replaces the old JSONB column for 10-100x faster similarity search. + embedding = Column(Vector(768), nullable=False) + embedding_model = Column(String(100), default="text-embedding-004") + + # Context metadata — powers context-aware answers with page citations + section_title = Column(String(500), nullable=True) # e.g. "Safety Procedures" + page_numbers = Column( + JSONB, nullable=True + ) # e.g. [12, 13] — pages this chunk spans + + # Legacy page columns (kept for backward compat, use page_numbers instead) + start_page = Column(Integer, nullable=True) + end_page = Column(Integer, nullable=True) + + extra_metadata = Column("metadata", JSONB, default=dict) + + # Relationships + document = relationship("Document", back_populates="embeddings") + + def __repr__(self): + return f"" diff --git a/app/models/prompt.py b/app/models/prompt.py new file mode 100644 index 0000000000000000000000000000000000000000..9363338fb77c4355433dd6f32efb636997148b92 --- /dev/null +++ b/app/models/prompt.py @@ -0,0 +1,63 @@ +""" +Custom Prompt Model +User-defined AI prompts for specialized mining document analysis +""" + +import uuid +from datetime import datetime + +from sqlalchemy import Boolean, Column, DateTime, ForeignKey, Integer, String, Text +from sqlalchemy.dialects.postgresql import UUID +from sqlalchemy.orm import relationship + +from app.models.base import Base, TimestampMixin, UUIDMixin + + +class CustomPrompt(Base, UUIDMixin, TimestampMixin): + """ + User-defined custom prompts for AI analysis. + Allows users to save specialized prompts for safety reviews, + compliance checks, equipment inspections, etc. + """ + + __tablename__ = "custom_prompts" + + # Owner + user_id = Column( + String(255), ForeignKey("users.clerk_user_id"), nullable=False, index=True + ) + + # Prompt info + name = Column(String(255), nullable=False) + prompt_text = Column(Text, nullable=False) + description = Column(Text, nullable=True) + + # Category/type for UI grouping + category = Column( + String(100), nullable=True + ) # e.g., "safety", "compliance", "equipment" + + # Whether this is a default/system prompt + is_default = Column(Boolean, default=False, nullable=False) + + # Usage tracking + use_count = Column(Integer, default=0, nullable=False) + + # Relationships + user = relationship("User", back_populates="custom_prompts") + + def __repr__(self): + return f"" + + def to_dict(self): + """Convert to dictionary for API responses""" + return { + "id": str(self.id), + "name": self.name, + "prompt": self.prompt_text, + "description": self.description, + "category": self.category, + "is_default": self.is_default, + "created_at": self.created_at.isoformat() if self.created_at else None, + "updated_at": self.updated_at.isoformat() if self.updated_at else None, + } diff --git a/app/models/user.py b/app/models/user.py new file mode 100644 index 0000000000000000000000000000000000000000..5e5ce9a3cef4185261d49ae10b829455e82b53c8 --- /dev/null +++ b/app/models/user.py @@ -0,0 +1,78 @@ +""" +User Model +User profile and organization management +""" + +import uuid +from datetime import datetime + +from sqlalchemy import JSON, Boolean, Column, DateTime, String, Text +from sqlalchemy.dialects.postgresql import JSONB as PG_JSONB +from sqlalchemy.dialects.postgresql import UUID + +JSONB = JSON().with_variant(PG_JSONB, "postgresql") +from sqlalchemy.orm import relationship + +from app.models.base import Base, TimestampMixin, UUIDMixin + + +class User(Base, UUIDMixin, TimestampMixin): + """ + User model linked to Clerk authentication. + Stores user profile and preferences. + """ + + __tablename__ = "users" + + # Clerk integration + clerk_user_id = Column(String(255), unique=True, nullable=False, index=True) + email = Column(String(255), nullable=True, index=True) + + # Profile + full_name = Column(String(255), nullable=True) + avatar_url = Column(Text, nullable=True) + + # Organization/Company (for enterprise) + company_name = Column(String(255), nullable=True) + company_role = Column( + String(100), nullable=True + ) # Safety Officer, Engineer, Manager + + # Mining-specific + industry_focus = Column(JSONB, nullable=True) # ["coal", "underground", "surface"] + mine_sites = Column(JSONB, nullable=True) # Associated mine sites + + # Preferences + preferences = Column(JSONB, default={}) + + # Status + is_active = Column(Boolean, default=True) + last_login = Column(DateTime, nullable=True) + + # Relationships + documents = relationship( + "Document", back_populates="user", cascade="all, delete-orphan" + ) + chat_sessions = relationship( + "ChatSession", back_populates="user", cascade="all, delete-orphan" + ) + custom_prompts = relationship( + "CustomPrompt", back_populates="user", cascade="all, delete-orphan" + ) + + def __repr__(self): + return f"" + + def to_dict(self): + """Convert to dictionary for API responses""" + return { + "id": str(self.id), + "clerk_user_id": self.clerk_user_id, + "email": self.email, + "full_name": self.full_name, + "company_name": self.company_name, + "company_role": self.company_role, + "industry_focus": self.industry_focus, + "is_active": self.is_active, + "created_at": self.created_at.isoformat() if self.created_at else None, + } diff --git a/app/schemas/__init__.py b/app/schemas/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..eaddd695ca09c7ebb8c598d49294c5dd4750cf38 --- /dev/null +++ b/app/schemas/__init__.py @@ -0,0 +1,50 @@ +""" +Pydantic Schemas +Request and response models for API validation +""" + +from app.schemas.analytics import DashboardStats, DocumentAnalytics, SafetyAnalytics +from app.schemas.chat import ( + ChatMessageResponse, + ChatRequest, + ChatResponse, + ChatSessionCreate, + ChatSessionResponse, +) +from app.schemas.common import ( + ErrorResponse, + HealthResponse, + JobStatusResponse, + PaginatedResponse, +) +from app.schemas.document import ( + DocumentAnalysisResponse, + DocumentCreate, + DocumentListResponse, + DocumentResponse, + DocumentUploadResponse, +) + +__all__ = [ + # Document + "DocumentCreate", + "DocumentResponse", + "DocumentListResponse", + "DocumentUploadResponse", + "DocumentAnalysisResponse", + # Chat + "ChatRequest", + "ChatResponse", + "ChatSessionCreate", + "ChatSessionResponse", + "ChatMessageResponse", + # Analytics + "DashboardStats", + "DocumentAnalytics", + "SafetyAnalytics", + # Common + "HealthResponse", + "ErrorResponse", + "PaginatedResponse", + "JobStatusResponse", +] diff --git a/app/schemas/analytics.py b/app/schemas/analytics.py new file mode 100644 index 0000000000000000000000000000000000000000..83f8b64a6465331866a578459af0e7869c67d365 --- /dev/null +++ b/app/schemas/analytics.py @@ -0,0 +1,140 @@ +""" +Analytics Schemas +Pydantic models for dashboard and analytics endpoints +""" + +from datetime import date, datetime +from typing import Any, Dict, List, Optional + +from pydantic import BaseModel, Field + + +class CategoryCount(BaseModel): + """Count by document category""" + + category: str + count: int + percentage: float + + +class StatusCount(BaseModel): + """Count by document status""" + + status: str + count: int + + +class SafetyDistribution(BaseModel): + """Safety score distribution""" + + range: str # "0-25", "26-50", "51-75", "76-100" + count: int + percentage: float + + +class DashboardStats(BaseModel): + """Main dashboard statistics""" + + # Document stats + total_documents: int = 0 + processed_documents: int = 0 + pending_documents: int = 0 + failed_documents: int = 0 + + # Chat stats + total_chat_sessions: int = 0 + total_messages: int = 0 + + # Safety stats + average_safety_score: Optional[float] = None + documents_with_hazards: int = 0 + compliance_violations: int = 0 + compliance_warnings: int = 0 + + # Processing stats + documents_processed_today: int = 0 + documents_processed_this_week: int = 0 + + # Category breakdown + documents_by_category: List[CategoryCount] = [] + + # Recent activity + last_upload_at: Optional[datetime] = None + last_chat_at: Optional[datetime] = None + + +class DocumentAnalytics(BaseModel): + """Detailed document analytics""" + + # Time series + uploads_by_day: List[Dict[str, Any]] = [] # [{"date": "2024-01-15", "count": 5}] + processing_times: List[Dict[str, Any]] = ( + [] + ) # [{"date": "...", "avg_time_ms": 1500}] + + # Category distribution + by_category: List[CategoryCount] = [] + + # Status distribution + by_status: List[StatusCount] = [] + + # File type distribution + by_file_type: List[Dict[str, Any]] = [] + + # Top documents by views + top_documents: List[Dict[str, Any]] = [] + + +class SafetyAnalytics(BaseModel): + """Safety compliance analytics""" + + # Overall scores + average_safety_score: float = 0 + median_safety_score: Optional[float] = None + min_safety_score: Optional[float] = None + max_safety_score: Optional[float] = None + + # Distribution + score_distribution: List[SafetyDistribution] = [] + + # Compliance + compliant_count: int = 0 + warning_count: int = 0 + violation_count: int = 0 + + # Hazards + total_hazards_detected: int = 0 + hazards_by_type: List[Dict[str, Any]] = [] # [{"type": "fall hazard", "count": 10}] + + # Trend + safety_trend: List[Dict[str, Any]] = [] # [{"date": "...", "avg_score": 75}] + + # Top issues + top_safety_concerns: List[Dict[str, Any]] = [] + + +class EntityAnalytics(BaseModel): + """Named entity analytics""" + + # Equipment mentioned + top_equipment: List[Dict[str, Any]] = ( + [] + ) # [{"name": "Caterpillar D11", "mentions": 25}] + + # Locations + top_locations: List[Dict[str, Any]] = [] + + # Chemicals + chemicals_mentioned: List[Dict[str, Any]] = [] + + # Regulations referenced + regulations_cited: List[Dict[str, Any]] = [] + + +class AnalyticsSummary(BaseModel): + """Combined analytics summary""" + + documents: DocumentAnalytics + safety: SafetyAnalytics + entities: EntityAnalytics + generated_at: datetime = Field(default_factory=datetime.utcnow) diff --git a/app/schemas/chat.py b/app/schemas/chat.py new file mode 100644 index 0000000000000000000000000000000000000000..5795002eda7af5ba283041e7a8c01e7addc66891 --- /dev/null +++ b/app/schemas/chat.py @@ -0,0 +1,113 @@ +""" +Chat Schemas +Pydantic models for chat API requests and responses +""" + +from datetime import datetime +from typing import Any, Dict, List, Optional + +from pydantic import BaseModel, Field + + +class ChatSource(BaseModel): + """Source citation for RAG response — includes page numbers for verifiable context.""" + + document_id: str + document_title: str + file_name: str # e.g. "Mining_site.pdf" + chunk_text: str = Field(..., max_length=500) + relevance_score: float = Field(..., ge=0, le=1) + page_numbers: List[int] = Field(default_factory=list) # e.g. [12, 13] + section_title: Optional[str] = None # e.g. "Safety Procedures" + + +class ChatRequest(BaseModel): + """Request to send a chat message""" + + content: str = Field(..., min_length=1, max_length=10000) + session_id: Optional[str] = Field( + None, description="Existing session ID or None for new session" + ) + document_ids: Optional[List[str]] = Field( + None, description="Specific documents to search" + ) + include_sources: bool = Field(default=True, description="Include source citations") + + +class ChatMessageResponse(BaseModel): + """Single chat message response""" + + id: str + role: str # "user" or "assistant" + content: str + sources: List[ChatSource] = [] + created_at: datetime + + # AI metadata (for assistant messages) + model_used: Optional[str] = None + response_time_ms: Optional[int] = None + + class Config: + from_attributes = True + + +class ChatResponse(BaseModel): + """Response after sending a message""" + + message: ChatMessageResponse + session_id: str + session_title: str + + +class ChatSessionCreate(BaseModel): + """Request to create a new chat session""" + + title: Optional[str] = Field("New Chat", max_length=200) + document_ids: Optional[List[str]] = Field( + None, description="Document context for session" + ) + system_prompt: Optional[str] = Field(None, description="Custom system prompt") + + +class ChatSessionResponse(BaseModel): + """Chat session response""" + + id: str + title: str + message_count: int + document_context: List[str] = [] + created_at: datetime + updated_at: Optional[datetime] = None + + # Preview of last message + last_message: Optional[str] = None + last_message_at: Optional[datetime] = None + + class Config: + from_attributes = True + + +class ChatSessionDetailResponse(BaseModel): + """Full chat session with messages""" + + id: str + title: str + document_context: List[str] = [] + system_prompt: Optional[str] = None + messages: List[ChatMessageResponse] = [] + created_at: datetime + updated_at: Optional[datetime] = None + + +class ChatSessionUpdateRequest(BaseModel): + """Request to update chat session""" + + title: Optional[str] = Field(None, max_length=200) + document_ids: Optional[List[str]] = None + + +class SuggestedQuestion(BaseModel): + """AI-suggested follow-up question""" + + question: str + category: Optional[str] = None # safety, equipment, regulatory diff --git a/app/schemas/common.py b/app/schemas/common.py new file mode 100644 index 0000000000000000000000000000000000000000..4105e46351d9306a6d7cc3aa70375abfe39ffe56 --- /dev/null +++ b/app/schemas/common.py @@ -0,0 +1,63 @@ +""" +Common Schemas +Shared response models and utilities +""" + +from datetime import datetime +from typing import Any, Generic, List, Optional, TypeVar + +from pydantic import BaseModel, Field + +T = TypeVar("T") + + +class HealthResponse(BaseModel): + """Health check response""" + + status: str = "healthy" + version: str + environment: str + timestamp: datetime = Field(default_factory=datetime.utcnow) + services: dict = Field(default_factory=dict) + + +class ErrorResponse(BaseModel): + """Standard error response""" + + error: str + code: str + details: Optional[dict] = None + timestamp: datetime = Field(default_factory=datetime.utcnow) + + +class PaginatedResponse(BaseModel, Generic[T]): + """Generic paginated response wrapper""" + + items: List[T] + total: int + page: int = 1 + page_size: int = 20 + total_pages: int + has_next: bool + has_prev: bool + + +class JobStatusResponse(BaseModel): + """Background job status response""" + + job_id: str + status: str # pending, processing, completed, failed + progress: Optional[int] = None # 0-100 + result: Optional[Any] = None + error: Optional[str] = None + created_at: datetime + updated_at: Optional[datetime] = None + completed_at: Optional[datetime] = None + + +class SuccessResponse(BaseModel): + """Generic success response""" + + success: bool = True + message: str = "Operation completed successfully" + data: Optional[Any] = None diff --git a/app/schemas/compliance.py b/app/schemas/compliance.py new file mode 100644 index 0000000000000000000000000000000000000000..f2f1612779ddd2b7e7bb4f532492185331c18996 --- /dev/null +++ b/app/schemas/compliance.py @@ -0,0 +1,68 @@ +""" +Compliance Audit Schemas +Pydantic v2 models for compliance audit API request/response validation. +""" + +from datetime import datetime +from typing import List, Optional +from uuid import UUID + +from pydantic import BaseModel, Field + + +class ComplianceAuditCreate(BaseModel): + """Request body for creating a new compliance audit.""" + + title: str = Field(..., min_length=1, max_length=500) + regulation_doc_id: UUID + operational_doc_ids: List[UUID] = Field(..., min_length=1) + + +class ComplianceMatrixRowResponse(BaseModel): + """A single clause assessment in the compliance matrix.""" + + id: UUID + clause_index: int + clause_text: str + section_title: Optional[str] = None + status: str # compliant | gap | missing + assessment: str + confidence: float + evidence_chunks: Optional[list] = None + recommendations: Optional[list] = None + + model_config = {"from_attributes": True} + + +class ComplianceAuditResponse(BaseModel): + """Compliance audit summary (without matrix rows).""" + + id: UUID + title: str + regulation_doc_id: UUID + operational_doc_ids: Optional[list] = None + status: str + total_clauses: Optional[int] = None + processed_clauses: int = 0 + compliant_count: Optional[int] = None + gap_count: Optional[int] = None + missing_count: Optional[int] = None + overall_score: Optional[float] = None + processing_error: Optional[str] = None + completed_at: Optional[datetime] = None + created_at: Optional[datetime] = None + + model_config = {"from_attributes": True} + + +class ComplianceAuditDetailResponse(ComplianceAuditResponse): + """Full audit detail including all matrix rows.""" + + rows: List[ComplianceMatrixRowResponse] = [] + + +class ComplianceAuditListResponse(BaseModel): + """Paginated list of compliance audits.""" + + audits: List[ComplianceAuditResponse] + total: int diff --git a/app/schemas/document.py b/app/schemas/document.py new file mode 100644 index 0000000000000000000000000000000000000000..0e8e3c8c63530f8ebdd8423801186e01399c2a2a --- /dev/null +++ b/app/schemas/document.py @@ -0,0 +1,171 @@ +""" +Document Schemas +Pydantic models for document API requests and responses +""" + +from datetime import datetime +from enum import Enum +from typing import Any, Dict, List, Optional + +from pydantic import BaseModel, Field, HttpUrl + + +class DocumentCategory(str, Enum): + """Document category enum matching the model""" + + SAFETY_PROTOCOL = "safety_protocol" + EQUIPMENT_MANUAL = "equipment_manual" + REGULATORY = "regulatory" + INCIDENT_REPORT = "incident_report" + GEOLOGICAL = "geological" + ENVIRONMENTAL = "environmental" + TRAINING = "training" + PERMIT = "permit" + MAINTENANCE = "maintenance" + OTHER = "other" + + +class DocumentStatus(str, Enum): + """Document processing status""" + + PENDING = "pending" + PROCESSING = "processing" + ANALYZING = "analyzing" + COMPLETED = "completed" + FAILED = "failed" + + +class ComplianceStatus(str, Enum): + """Safety compliance status""" + + COMPLIANT = "compliant" + WARNING = "warning" + VIOLATION = "violation" + PENDING = "pending" + NOT_APPLICABLE = "not_applicable" + + +class DocumentCreate(BaseModel): + """Request model for creating a document from UploadThing""" + + file_url: str = Field(..., description="UploadThing file URL") + file_name: str = Field(..., description="Original file name") + file_size: int = Field(..., gt=0, description="File size in bytes") + file_type: str = Field(..., description="MIME type") + title: Optional[str] = Field(None, description="Custom document title") + tags: Optional[List[str]] = Field(default=[], description="User-defined tags") + + +class DocumentUploadResponse(BaseModel): + """Response after document upload - before processing""" + + id: str + title: str + file_name: str + status: DocumentStatus = DocumentStatus.PENDING + job_id: Optional[str] = Field(None, description="Background processing job ID") + message: str = "Document uploaded successfully. Processing started." + + +class DocumentAnalysisResult(BaseModel): + """AI analysis results for a document""" + + # Classification + category: DocumentCategory + subcategory: Optional[str] = None + classification_confidence: float = Field(..., ge=0, le=1) + + # Summary + summary: str + key_points: List[str] = [] + + # Safety Analysis + safety_score: Optional[float] = Field(None, ge=0, le=100) + compliance_status: ComplianceStatus = ComplianceStatus.PENDING + hazards_detected: List[Dict[str, Any]] = [] + safety_recommendations: List[str] = [] + reasoning: Optional[Dict[str, Any]] = None + + # Entities + entities: Dict[str, List[str]] = Field( + default_factory=lambda: { + "equipment": [], + "chemicals": [], + "locations": [], + "personnel": [], + "dates": [], + "regulations": [], + } + ) + + +class DocumentResponse(BaseModel): + """Full document response with all fields""" + + id: str + title: str + file_name: str + file_size: int + file_type: str + file_url: str + status: DocumentStatus + + # Processing info + processing_error: Optional[str] = None + processed_at: Optional[datetime] = None + + # Content info + page_count: Optional[int] = None + word_count: Optional[int] = None + + # AI Classification + category: Optional[DocumentCategory] = None + subcategory: Optional[str] = None + classification_confidence: Optional[float] = None + + # AI Summary + summary: Optional[str] = None + key_points: Optional[List[str]] = None + + # Safety Analysis + safety_score: Optional[float] = None + compliance_status: Optional[ComplianceStatus] = None + hazards_detected: Optional[List[Dict[str, Any]]] = None + safety_recommendations: Optional[List[str]] = None + + # Entities + entities: Optional[Dict[str, List[str]]] = None + + # Metadata + tags: List[str] = [] + created_at: datetime + + class Config: + from_attributes = True + + +class DocumentListResponse(BaseModel): + """Paginated list of documents""" + + documents: List[DocumentResponse] + total: int + page: int = 1 + page_size: int = 20 + + # Aggregated stats + stats: Optional[Dict[str, Any]] = Field( + default_factory=lambda: { + "by_category": {}, + "by_status": {}, + "avg_safety_score": None, + } + ) + + +class DocumentAnalysisResponse(BaseModel): + """Response for document analysis endpoint""" + + document_id: str + status: str + analysis: Optional[DocumentAnalysisResult] = None + processing_time_ms: Optional[int] = None diff --git a/app/schemas/prompt.py b/app/schemas/prompt.py new file mode 100644 index 0000000000000000000000000000000000000000..b4155cc5fb51f7c0d411046592ddbe0d87ff1876 --- /dev/null +++ b/app/schemas/prompt.py @@ -0,0 +1,52 @@ +""" +Prompt Schemas +Pydantic models for custom prompt API requests and responses +""" + +from datetime import datetime +from typing import List, Optional + +from pydantic import BaseModel, Field + + +class PromptCreate(BaseModel): + """Request model for creating a custom prompt""" + + name: str = Field(..., min_length=1, max_length=255, description="Prompt name") + prompt: str = Field(..., min_length=1, description="The prompt text") + description: Optional[str] = Field(None, description="Optional description") + category: Optional[str] = Field( + None, description="Category for grouping (e.g. safety, compliance)" + ) + + +class PromptUpdate(BaseModel): + """Request model for updating a custom prompt""" + + name: Optional[str] = Field(None, min_length=1, max_length=255) + prompt: Optional[str] = Field(None, min_length=1) + description: Optional[str] = None + category: Optional[str] = None + + +class PromptResponse(BaseModel): + """Response model for a custom prompt""" + + id: str + name: str + prompt: str + description: Optional[str] = None + category: Optional[str] = None + is_default: bool = False + created_at: datetime + updated_at: Optional[datetime] = None + + class Config: + from_attributes = True + + +class PromptListResponse(BaseModel): + """Response model for listing prompts""" + + prompts: List[PromptResponse] + total: int diff --git a/app/services/__init__.py b/app/services/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..cca0d1dfb3a2f76fca5feb347e1c2467db2776eb --- /dev/null +++ b/app/services/__init__.py @@ -0,0 +1,13 @@ +""" +Services Module +Business logic layer +""" + +from app.services.chat_service import ChatService +from app.services.document_service import DocumentService, process_document_async + +__all__ = [ + "DocumentService", + "ChatService", + "process_document_async", +] diff --git a/app/services/chat_service.py b/app/services/chat_service.py new file mode 100644 index 0000000000000000000000000000000000000000..69a3b8be7325094d90b2276ddcab8241de54e169 --- /dev/null +++ b/app/services/chat_service.py @@ -0,0 +1,310 @@ +""" +Chat Service +RAG-powered conversation with production-grade retrieval pipeline. + +Pipeline: Query → Embed → Hybrid Search (Vector + BM25) → Rerank → LLM + +Key improvements over v1: + - Hybrid search: pgvector cosine + pg_trgm BM25 via Reciprocal Rank Fusion + - Cross-encoder reranking: ms-marco-MiniLM-L-6-v2 for precise relevance + - Similarity threshold: filters irrelevant chunks before context formatting + - System prompt as system role: proper LLM message structure + - Context includes page numbers, section titles, and file names + - Streaming support via generate_response_stream() +""" + +import asyncio +import json +import logging +from typing import Any, AsyncGenerator, Dict, List, Optional, Tuple + +import google.generativeai as genai +from sqlalchemy.orm import Session + +from app.config import settings + +logger = logging.getLogger(__name__) + +# Configure Gemini once +genai.configure(api_key=settings.GEMINI_API_KEY) + +# ── System Prompt ────────────────────────────────────────────────────────────── +_SYSTEM_PROMPT = """You are MiningNiti AI, an expert assistant specialized in the coal mining industry. + +## CRITICAL RULES +1. ONLY cite documents that appear in the "Document Context" section below. + NEVER invent, guess, or fabricate document names that are not provided in the context. +2. ONLY use the exact file names and page numbers shown in the context chunks. +3. If the context does not contain enough information to answer, say: + "Based on the available documents, I cannot find specific information about this topic." +4. NEVER make up facts, regulations, or details not found in the provided context. + +## Citation Format +Cite using the exact file name and page from the context: + [FileName.pdf, Page X] + +Example (if the context shows "Source: acts_1948.pdf, Page: 5"): + "As per [acts_1948.pdf, Page 5], the act provides for..." + +If a claim spans multiple pages: + [acts_1948.pdf, Pages 1-3] + +## Behavior Rules +1. ALWAYS prioritize safety information +2. ALWAYS cite the source document and page number for factual claims +3. Provide practical, actionable guidance +4. If the document context doesn't answer the question, clearly say so +5. Reference specific regulations when applicable +""" + + +class ChatService: + """ + Chat service with production RAG pipeline. + + Retrieval flow: + 1. Embed query via Gemini text-embedding-004 + 2. Hybrid search: pgvector cosine similarity + pg_trgm BM25 + 3. Reciprocal Rank Fusion to merge results + 4. Cross-encoder reranking for precise relevance scoring + 5. Top-K chunks formatted as context for LLM + 6. System prompt as system role for proper instruction following + """ + + def __init__(self): + self.model = genai.GenerativeModel(settings.GEMINI_MODEL) + self.embedding_model = settings.EMBEDDING_MODEL + + # ── Public API ───────────────────────────────────────────────────────────── + + async def generate_response( + self, + query: str, + user_id: str, + document_ids: Optional[List[str]] = None, + db: Session = None, + ) -> Tuple[str, List[Dict[str, Any]]]: + try: + query_embedding = await self._get_embedding(query) + relevant_chunks = await self._retrieve_chunks( + query=query, + query_embedding=query_embedding, + user_id=user_id, + document_ids=document_ids, + db=db, + ) + + context = self._format_context(relevant_chunks) + + from app.services.llm_provider import get_groq_client + + client = get_groq_client() + + response = await client.chat.completions.create( + model="llama-3.3-70b-versatile", + messages=[ + {"role": "system", "content": _SYSTEM_PROMPT}, + { + "role": "user", + "content": self._build_user_message(query, context), + }, + ], + ) + answer = response.choices[0].message.content + + tokens_used = ( + { + "input": response.usage.prompt_tokens, + "output": response.usage.completion_tokens, + } + if hasattr(response, "usage") + else None + ) + + sources = self._build_sources(relevant_chunks[:3]) + return answer, sources, tokens_used + + except Exception as e: + logger.error(f"Chat generation error: {e}", exc_info=True) + return ( + "I apologize, but I encountered an error processing your question. Please try again.", + [], + None, + ) + + async def generate_response_stream( + self, + query: str, + user_id: str, + document_ids: Optional[List[str]] = None, + db: Session = None, + ) -> AsyncGenerator[str, None]: + try: + query_embedding = await self._get_embedding(query) + relevant_chunks = await self._retrieve_chunks( + query=query, + query_embedding=query_embedding, + user_id=user_id, + document_ids=document_ids, + db=db, + ) + + sources = self._build_sources(relevant_chunks[:3]) + yield f"event: sources\ndata: {json.dumps(sources)}\n\n" + + context = self._format_context(relevant_chunks) + + from app.services.llm_provider import get_groq_client + + client = get_groq_client() + + response_stream = await client.chat.completions.create( + model="llama-3.3-70b-versatile", + messages=[ + {"role": "system", "content": _SYSTEM_PROMPT}, + { + "role": "user", + "content": self._build_user_message(query, context), + }, + ], + stream=True, + ) + + async for chunk in response_stream: + if chunk.choices[0].delta.content: + yield f"event: token\ndata: {json.dumps({'text': chunk.choices[0].delta.content})}\n\n" + + tokens_used = None + yield f"event: done\ndata: {json.dumps({'sources_count': len(sources), 'tokens_used': tokens_used})}\n\n" + + except Exception as e: + logger.error(f"Stream generation error: {e}", exc_info=True) + yield f"event: error\ndata: {json.dumps({'message': str(e)})}\n\n" + + def get_mining_suggestions(self) -> List[str]: + """Get suggested mining-related questions.""" + return [ + "What are the MSHA requirements for underground coal mines?", + "How often should mining equipment be inspected?", + "What are the emergency evacuation procedures?", + "What are the ventilation requirements for coal mines?", + "How to conduct a proper safety audit?", + ] + + # ── Retrieval Pipeline ──────────────────────────────────────────────────── + + async def _retrieve_chunks( + self, + query: str, + query_embedding: List[float], + user_id: str, + document_ids: Optional[List[str]], + db: Session, + ) -> List[Dict[str, Any]]: + """ + Production RAG retrieval pipeline: + 1. Hybrid search (vector + BM25) → over-fetch candidates + 2. Cross-encoder rerank → precise top-K + """ + if not query_embedding: + return [] + + # Step 1: Hybrid search (over-fetch for reranking) + from app.services.hybrid_search import hybrid_search + + candidates = await hybrid_search( + query_text=query, + query_embedding=query_embedding, + db=db, + user_id=user_id, + document_ids=document_ids, + top_k=settings.RERANK_OVER_FETCH, + ) + + if not candidates: + return [] + + # Step 2: Cross-encoder reranking + if settings.ENABLE_RERANKING and len(candidates) > settings.RERANK_TOP_K: + from app.services.reranker import rerank + + candidates = rerank( + query=query, + chunks=candidates, + top_k=settings.RERANK_TOP_K, + ) + + return candidates + + # ── Embedding ───────────────────────────────────────────────────────────── + + async def _get_embedding(self, text: str) -> List[float]: + """Get embedding for query text using Gemini.""" + try: + result = await asyncio.to_thread( + genai.embed_content, + model=self.embedding_model, + content=text, + task_type="retrieval_query", + output_dimensionality=768, + ) + return result["embedding"] + except Exception as e: + logger.error(f"Embedding generation failed: {e}") + return [] + + # ── Context Formatting ──────────────────────────────────────────────────── + + def _format_context(self, chunks: List[Dict[str, Any]]) -> str: + if not chunks: + return "No relevant document context found." + + context_parts = [] + for i, chunk in enumerate(chunks, 1): + pages = chunk.get("page_numbers", []) + page_str = ( + f"{pages[0]}-{pages[-1]}" + if len(pages) > 1 + else f"{pages[0]}" if pages else "unknown" + ) + file_name = chunk.get("file_name", "unknown.pdf") + text = chunk.get("text", "") + section = chunk.get("section_title", "") + + section_prefix = f"[{section}] " if section else "" + context_parts.append( + f"Context chunk {i}: {section_prefix}{text} (Source: {file_name}, Page: {page_str})" + ) + + return "\n".join(context_parts) + + def _build_user_message(self, query: str, context: str) -> str: + """Build the user message with context and query.""" + return ( + f"## Document Context:\n{context}\n\n" + f"## User Question:\n{query}\n\n" + f"## Your Answer (remember to cite [FileName, Page X] for every claim):\n" + ) + + def _build_sources(self, chunks: List[Dict[str, Any]]) -> List[Dict[str, Any]]: + """Build structured source citations for API response.""" + sources = [] + for chunk in chunks: + pages = chunk.get("page_numbers", []) + sources.append( + { + "document_id": chunk["document_id"], + "document_title": chunk["document_title"], + "file_name": chunk["file_name"], + "file_url": chunk.get("file_url"), + "chunk_text": chunk["text"][:300] + + ("..." if len(chunk["text"]) > 300 else ""), + "exact_text_chunk": chunk["text"], + "relevance_score": round( + chunk.get("rerank_score", chunk.get("score", 0.0)), 4 + ), + "page_numbers": pages, + "section_title": chunk.get("section_title"), + } + ) + return sources diff --git a/app/services/chunking.py b/app/services/chunking.py new file mode 100644 index 0000000000000000000000000000000000000000..49135cf4aa7ba164ce3bd9c277416ddad5598551 --- /dev/null +++ b/app/services/chunking.py @@ -0,0 +1,290 @@ +""" +Smart Document Chunking Service +Sentence-aware text chunking with page number and section tracking. + +Replaces the crude word-count split (text.split()) with: + - Sentence-boundary detection (no mid-sentence breaks) + - Per-chunk page number list (e.g. [12, 13]) + - Section/heading detection from document structure + - Configurable overlap to preserve cross-boundary context +""" + +import logging +import re +from dataclasses import dataclass, field +from typing import List, Optional + +from app.config import settings +from app.services.extractors import PageContent + +logger = logging.getLogger(__name__) + + +# ── Data structures ──────────────────────────────────────────────────────────── + + +@dataclass +class DocumentChunk: + """A single text chunk with full provenance metadata.""" + + chunk_index: int + text: str + page_numbers: List[int] # Pages this chunk spans, e.g. [12, 13] + section_title: Optional[str] # Nearest detected heading, e.g. "Safety Procedures" + char_start: int = 0 # Character offset in full_text + char_end: int = 0 + + +# ── Heading detection patterns ───────────────────────────────────────────────── + +# Matches ALL-CAPS lines (≥4 chars), numbered sections (1.2.3), or markdown headings +_HEADING_PATTERNS = [ + re.compile(r"^#{1,4}\s+(.+)$", re.MULTILINE), # Markdown headings + re.compile( + r"^(\d+(?:\.\d+)*)\s+([A-Z][^\n]{3,60})$", re.MULTILINE + ), # Numbered: "1.2 Section" + re.compile(r"^([A-Z][A-Z\s\-]{4,60})$", re.MULTILINE), # ALL-CAPS headings +] + +# Sentence boundary: period/question/exclamation followed by space and capital (or end) +_SENTENCE_BOUNDARY = re.compile(r"(?<=[.!?])\s+(?=[A-Z\"])") + + +class ChunkingService: + """ + Sentence-aware document chunker with page number tracking. + + Algorithm: + 1. Split full text into sentences using regex boundary detection + 2. Group sentences into chunks respecting max token size + 3. Track which pages each chunk spans using char offsets + 4. Detect section headings and annotate each chunk with nearest heading + 5. Add configurable word overlap between adjacent chunks + """ + + def __init__( + self, + chunk_size: int = None, # words per chunk + chunk_overlap: int = None, # words of overlap + min_chunk_words: int = 20, # skip chunks smaller than this + ): + self.chunk_size = chunk_size or settings.CHUNK_SIZE # default: 1000 + self.chunk_overlap = chunk_overlap or settings.CHUNK_OVERLAP # default: 200 + self.min_chunk_words = min_chunk_words + + # ── Public API ───────────────────────────────────────────────────────────── + + def chunk_document( + self, + full_text: str, + pages: List[PageContent], + ) -> List[DocumentChunk]: + """ + Split a document into annotated chunks. + + Args: + full_text: Complete document text + pages: Per-page content from extractor (for page number mapping) + + Returns: + List of DocumentChunk objects with page_numbers and section_title + """ + if not full_text.strip(): + return [] + + # Build page offset map: char_offset → page_number + page_map = self._build_page_map(pages) + + # Extract section headings with their char positions + headings = self._extract_headings(full_text) + + # Split into sentences + sentences = self._split_sentences(full_text) + + if not sentences: + return [] + + # Group sentences into word-count-bounded chunks with overlap + raw_chunks = self._group_into_chunks(sentences, full_text) + + # Annotate each chunk with page numbers and section title + chunks: List[DocumentChunk] = [] + for idx, (chunk_text, char_start, char_end) in enumerate(raw_chunks): + if len(chunk_text.split()) < self.min_chunk_words: + continue # Skip tiny fragments + + page_nums = self._get_page_numbers(char_start, char_end, page_map) + section = self._get_section_title(char_start, headings) + + chunks.append( + DocumentChunk( + chunk_index=idx, + text=chunk_text.strip(), + page_numbers=page_nums, + section_title=section, + char_start=char_start, + char_end=char_end, + ) + ) + + logger.info( + f"Chunked document: {len(sentences)} sentences → {len(chunks)} chunks " + f"(size={self.chunk_size} words, overlap={self.chunk_overlap} words)" + ) + return chunks + + # ── Internal helpers ─────────────────────────────────────────────────────── + + def _split_sentences(self, text: str) -> List[str]: + """Split text into sentences using regex boundary detection.""" + # Replace common abbreviations that fool period detection + text = re.sub( + r"\b(Mr|Mrs|Ms|Dr|Prof|Sr|Jr|vs|etc|No|Vol|Fig)\.", r"\1", text + ) + + # Split on sentence boundaries + raw_sentences = _SENTENCE_BOUNDARY.split(text) + + # Restore abbreviation dots + sentences = [ + s.replace("", ".").strip() for s in raw_sentences if s.strip() + ] + return sentences + + def _group_into_chunks( + self, + sentences: List[str], + full_text: str, + ) -> List[tuple]: + """ + Group sentences into chunks respecting chunk_size with overlap. + + Returns list of (chunk_text, char_start, char_end) tuples. + """ + chunks = [] + current_sentences: List[str] = [] + current_word_count = 0 + + # Precompute sentence char offsets in full_text + sentence_offsets = self._compute_sentence_offsets(sentences, full_text) + + i = 0 + while i < len(sentences): + sentence = sentences[i] + word_count = len(sentence.split()) + + if current_word_count + word_count <= self.chunk_size: + current_sentences.append(sentence) + current_word_count += word_count + i += 1 + else: + if current_sentences: + # Save current chunk + start_idx = sentences.index(current_sentences[0]) + end_idx = sentences.index(current_sentences[-1]) + char_start = sentence_offsets[start_idx][0] + char_end = sentence_offsets[end_idx][1] + chunks.append((" ".join(current_sentences), char_start, char_end)) + + # Build overlap: keep last N words worth of sentences + overlap_sentences = self._get_overlap_sentences( + current_sentences, self.chunk_overlap + ) + current_sentences = overlap_sentences + current_word_count = sum(len(s.split()) for s in current_sentences) + else: + # Single sentence is longer than chunk_size — add it anyway + char_start, char_end = sentence_offsets[i] + chunks.append((sentence, char_start, char_end)) + i += 1 + + # Don't forget the last chunk + if current_sentences: + start_idx = sentences.index(current_sentences[0]) + end_idx = sentences.index(current_sentences[-1]) + char_start = sentence_offsets[start_idx][0] + char_end = sentence_offsets[end_idx][1] + chunks.append((" ".join(current_sentences), char_start, char_end)) + + return chunks + + def _get_overlap_sentences( + self, sentences: List[str], target_words: int + ) -> List[str]: + """Return the tail sentences that total approximately target_words words.""" + result = [] + word_count = 0 + for sentence in reversed(sentences): + wc = len(sentence.split()) + if word_count + wc > target_words: + break + result.insert(0, sentence) + word_count += wc + return result + + def _compute_sentence_offsets( + self, sentences: List[str], full_text: str + ) -> List[tuple]: + """Find (start, end) char offsets of each sentence in full_text.""" + offsets = [] + search_from = 0 + for sentence in sentences: + # Find the sentence in full text starting from last known position + idx = full_text.find(sentence[:30], search_from) # match on first 30 chars + if idx == -1: + idx = search_from + end = idx + len(sentence) + offsets.append((idx, end)) + search_from = max(search_from, idx + 1) + return offsets + + def _build_page_map(self, pages: List[PageContent]) -> List[tuple]: + """Build sorted list of (char_start, char_end, page_number) for binary search.""" + return [ + (p.char_start, p.char_end, p.page_number) + for p in sorted(pages, key=lambda p: p.char_start) + ] + + def _get_page_numbers( + self, char_start: int, char_end: int, page_map: List[tuple] + ) -> List[int]: + """Return all page numbers that a chunk's char range overlaps.""" + page_nums = [] + for p_start, p_end, page_num in page_map: + # Overlap condition + if p_start < char_end and p_end > char_start: + page_nums.append(page_num) + return sorted(set(page_nums)) or [1] + + def _extract_headings(self, text: str) -> List[tuple]: + """ + Extract (char_position, heading_text) from document. + Looks for ALL-CAPS lines, numbered sections, and markdown headings. + """ + headings = [] + for pattern in _HEADING_PATTERNS: + for match in pattern.finditer(text): + heading_text = match.group(0).strip() + # Clean up heading text + heading_text = re.sub(r"^#+\s*", "", heading_text) # Remove markdown # + heading_text = re.sub( + r"^\d+(?:\.\d+)*\s*", "", heading_text + ) # Remove numbering + if 3 <= len(heading_text) <= 100: + headings.append((match.start(), heading_text)) + + # Sort by position + headings.sort(key=lambda h: h[0]) + return headings + + def _get_section_title( + self, char_start: int, headings: List[tuple] + ) -> Optional[str]: + """Return the most recent heading before char_start.""" + result = None + for pos, title in headings: + if pos <= char_start: + result = title + else: + break + return result diff --git a/app/services/compliance_service.py b/app/services/compliance_service.py new file mode 100644 index 0000000000000000000000000000000000000000..036f7892ff2863718358d46428b4bd7bde115a41 --- /dev/null +++ b/app/services/compliance_service.py @@ -0,0 +1,341 @@ +""" +Compliance Service +Runs compliance audits by cross-referencing regulation document clauses +against operational document evidence using pgvector similarity search +and the ComplianceAuditorAgent. + +Pipeline: + 1. Load audit record + regulation doc + operational docs + 2. Extract regulation clauses from regulation doc embeddings + 3. For each clause: pgvector cosine search → top-K evidence chunks + 4. ComplianceAuditorAgent assesses each clause (parallel, capped) + 5. Persist results, compute aggregate stats, mark audit COMPLETED +""" + +import asyncio +import logging +from datetime import datetime, timezone +from typing import Dict, List, Optional + +from sqlalchemy import text +from sqlalchemy.orm import Session + +from app.agents.base import QuotaExceededError +from app.agents.compliance_auditor import ComplianceAuditorAgent +from app.config import settings +from app.db.session import get_db_context +from app.models.compliance import AuditStatus, ComplianceAudit, ComplianceMatrixRow +from app.models.document import Document, DocumentEmbedding + +logger = logging.getLogger(__name__) + +# Max concurrent LLM calls to avoid rate limits +_MAX_CONCURRENT_ASSESSMENTS = 5 +# Top-K evidence chunks per clause (after reranking) +_TOP_K = 5 +# Over-fetch count for reranking +_OVER_FETCH = 15 +# Relevance threshold for evidence chunks +_RELEVANCE_THRESHOLD = 0.30 + + +class ComplianceService: + """ + Compliance audit execution service. + + Uses pgvector cosine similarity to find operational document evidence + for each regulation clause, then runs ComplianceAuditorAgent assessments. + """ + + def __init__(self): + self.agent = ComplianceAuditorAgent() + self._semaphore = asyncio.Semaphore(_MAX_CONCURRENT_ASSESSMENTS) + + async def run_audit(self, audit_id: str) -> bool: + """ + Execute a full compliance audit. + + Returns True on success, False on failure. + """ + logger.info(f"Starting compliance audit: {audit_id}") + + with get_db_context() as db: + audit = ( + db.query(ComplianceAudit).filter(ComplianceAudit.id == audit_id).first() + ) + + if not audit: + logger.error(f"Audit not found: {audit_id}") + return False + + try: + # ── Step 1: Mark as running ───────────────────────────── + audit.status = AuditStatus.RUNNING + db.commit() + + # ── Step 2: Load regulation doc embeddings (clauses) ──── + reg_doc = ( + db.query(Document) + .filter(Document.id == audit.regulation_doc_id) + .first() + ) + if not reg_doc or reg_doc.status.value != "completed": + raise ValueError( + f"Regulation document not ready: {audit.regulation_doc_id}" + ) + + reg_embeddings = ( + db.query(DocumentEmbedding) + .filter(DocumentEmbedding.document_id == reg_doc.id) + .order_by(DocumentEmbedding.chunk_index) + .all() + ) + if not reg_embeddings: + raise ValueError( + "Regulation document has no embeddings. " + "Ensure it was fully processed before auditing." + ) + + audit.total_clauses = len(reg_embeddings) + db.commit() + + logger.info( + f"Audit {audit_id}: {len(reg_embeddings)} regulation clauses " + f"to assess against {len(audit.operational_doc_ids)} operational docs" + ) + + # ── Step 3: Assess each clause ────────────────────────── + operational_doc_ids = [ + str(d) for d in (audit.operational_doc_ids or []) + ] + + results = await self._assess_clauses( + db=db, + audit_id=audit_id, + user_id=audit.user_id, + clauses=reg_embeddings, + operational_doc_ids=operational_doc_ids, + ) + + # ── Step 4: Persist matrix rows ───────────────────────── + compliant_count = 0 + gap_count = 0 + missing_count = 0 + + for i, result in enumerate(results): + status = result.get("status", "missing") + if status == "compliant": + compliant_count += 1 + elif status == "gap": + gap_count += 1 + else: + missing_count += 1 + + row = ComplianceMatrixRow( + audit_id=audit.id, + clause_index=i, + clause_text=result.get("clause_text", ""), + section_title=result.get("section_title"), + status=status, + assessment=result.get("assessment", ""), + confidence=result.get("confidence", 0.5), + evidence_chunks=result.get("evidence_chunks", []), + recommendations=result.get("recommendations", []), + ) + db.add(row) + + # ── Step 5: Compute aggregate stats ───────────────────── + total = len(results) + audit.compliant_count = compliant_count + audit.gap_count = gap_count + audit.missing_count = missing_count + audit.processed_clauses = total + audit.overall_score = round( + (compliant_count / total * 100) if total > 0 else 0, 1 + ) + + # ── Step 6: Mark completed ────────────────────────────── + audit.status = AuditStatus.COMPLETED + audit.completed_at = datetime.now(timezone.utc) + db.commit() + + logger.info( + f"Compliance audit completed: {audit_id} — " + f"Score: {audit.overall_score}%, " + f"Compliant: {compliant_count}, Gaps: {gap_count}, " + f"Missing: {missing_count}" + ) + return True + + except QuotaExceededError as qe: + logger.error(f"Quota exceeded during audit {audit_id}: {qe}") + audit.status = AuditStatus.COMPLETED + audit.processing_error = ( + "AI analysis incomplete: API quota exceeded. " + "Partial results saved. Re-run when quota resets." + ) + audit.completed_at = datetime.now(timezone.utc) + db.commit() + return False + + except Exception as e: + logger.error( + f"Compliance audit failed: {audit_id} — {e}", exc_info=True + ) + audit.status = AuditStatus.FAILED + audit.processing_error = str(e) + db.commit() + return False + + async def _assess_clauses( + self, + db: Session, + audit_id: str, + user_id: str, + clauses: List[DocumentEmbedding], + operational_doc_ids: List[str], + ) -> List[Dict]: + """ + Assess all clauses in parallel with semaphore-capped concurrency. + + Returns a list of result dicts (one per clause, in order). + """ + results: List[Optional[Dict]] = [None] * len(clauses) + + async def assess_one(index: int, clause: DocumentEmbedding): + async with self._semaphore: + try: + # Hybrid search for evidence + rerank + evidence = await self._find_evidence( + db=db, + user_id=user_id, + query_embedding=clause.embedding, + operational_doc_ids=operational_doc_ids, + clause_text=clause.chunk_text, + ) + + # Run agent assessment + result = await self.agent.analyze( + text=clause.chunk_text, + context={ + "evidence_chunks": evidence, + "clause_section": clause.section_title or "", + }, + ) + + results[index] = { + "clause_text": clause.chunk_text, + "section_title": clause.section_title, + "status": result.get("status", "missing"), + "assessment": result.get("assessment", ""), + "confidence": result.get("confidence", 0.5), + "evidence_chunks": evidence, + "recommendations": result.get("recommendations", []), + } + + # Update progress counter + with get_db_context() as progress_db: + progress_audit = ( + progress_db.query(ComplianceAudit) + .filter(ComplianceAudit.id == audit_id) + .first() + ) + if progress_audit: + progress_audit.processed_clauses = sum( + 1 for r in results if r is not None + ) + progress_db.commit() + + except Exception as e: + logger.error( + f"Clause {index} assessment failed: {e}", exc_info=True + ) + results[index] = { + "clause_text": clause.chunk_text, + "section_title": clause.section_title, + "status": "missing", + "assessment": f"Assessment failed: {str(e)}", + "confidence": 0.0, + "evidence_chunks": [], + "recommendations": [], + } + + # Launch all clause assessments + tasks = [assess_one(i, clause) for i, clause in enumerate(clauses)] + await asyncio.gather(*tasks, return_exceptions=True) + + # Fill any remaining None slots (shouldn't happen but safety net) + for i, r in enumerate(results): + if r is None: + results[i] = { + "clause_text": clauses[i].chunk_text, + "section_title": clauses[i].section_title, + "status": "missing", + "assessment": "Assessment did not complete.", + "confidence": 0.0, + "evidence_chunks": [], + "recommendations": [], + } + + return results + + async def _find_evidence( + self, + db: Session, + user_id: str, + query_embedding: list, + operational_doc_ids: List[str], + clause_text: str = "", + ) -> List[Dict]: + """ + Hybrid search for evidence chunks from operational documents. + + Pipeline: pgvector cosine + pg_trgm BM25 → RRF → cross-encoder rerank + """ + if not operational_doc_ids: + return [] + + # Step 1: Hybrid search (over-fetch for reranking) + from app.services.hybrid_search import hybrid_search + + candidates = await hybrid_search( + query_text=clause_text, + query_embedding=query_embedding, + db=db, + user_id=user_id, + document_ids=operational_doc_ids, + top_k=_OVER_FETCH, + ) + + # Step 2: Rerank if we have more candidates than needed + if settings.ENABLE_RERANKING and len(candidates) > _TOP_K: + from app.services.reranker import rerank + + candidates = rerank( + query=clause_text, + chunks=candidates, + top_k=_TOP_K, + text_key="text", + ) + else: + candidates = candidates[:_TOP_K] + + return [ + { + "chunk_text": c["text"], + "document_title": c["document_title"], + "page_numbers": c.get("page_numbers", []), + "section_title": c.get("section_title"), + "relevance_score": round(c.get("rerank_score", c.get("score", 0.0)), 4), + } + for c in candidates + ] + + +# ── Background task wrapper ──────────────────────────────────────────────────── + + +async def run_compliance_audit_async(audit_id: str) -> None: + """Async wrapper used with the compliance task queue.""" + service = ComplianceService() + await service.run_audit(audit_id) diff --git a/app/services/document_service.py b/app/services/document_service.py new file mode 100644 index 0000000000000000000000000000000000000000..69c0a8ec8d1137242f98985d4733777d970d555b --- /dev/null +++ b/app/services/document_service.py @@ -0,0 +1,241 @@ +""" +Document Service +Document processing pipeline with unified multi-agent AI analysis. + +Pipeline: + 1. Download file (with SSRF validation) + 2. Extract text + page boundaries via extractors.py + 3. Smart sentence-aware chunking via chunking.py + 4. Generate pgvector embeddings (stored as Vector(768)) + 5. Run all 4 agents via AgentOrchestrator (parallel execution) + 6. Persist results and mark document COMPLETED +""" + +import logging +from datetime import datetime, timezone +from typing import List + +import google.generativeai as genai + +from app.agents.base import QuotaExceededError +from app.agents.orchestrator import AgentOrchestrator +from app.config import settings +from app.db.session import get_db_context +from app.models.document import ( + ComplianceStatus, + Document, + DocumentCategory, + DocumentEmbedding, + DocumentStatus, +) +from app.services.chunking import ChunkingService, DocumentChunk +from app.services.extractors import ExtractedDocument, download_and_extract + +logger = logging.getLogger(__name__) + +# Configure Gemini once at module level +genai.configure(api_key=settings.GEMINI_API_KEY) + +_chunker = ChunkingService() +_orchestrator = AgentOrchestrator() + + +class DocumentService: + """ + Document processing service. + + Uses: + - extractors.py → page-aware text extraction + - chunking.py → sentence-aware chunking with page tracking + - orchestrator → unified 4-agent parallel analysis + - pgvector → native vector storage (no brute-force cosine in Python) + """ + + def __init__(self): + self.embedding_model = settings.EMBEDDING_MODEL + + async def process_document(self, document_id: str) -> bool: + """ + Full document processing pipeline. + + Returns True on success, False on failure. + """ + logger.info(f"Starting document processing: {document_id}") + + with get_db_context() as db: + document = db.query(Document).filter(Document.id == document_id).first() + + if not document: + logger.error(f"Document not found: {document_id}") + return False + + try: + # ── Step 1: Mark as processing ────────────────────────────── + document.status = DocumentStatus.PROCESSING + db.commit() + + # ── Step 2: Download and extract text with page tracking ──── + logger.info(f"Extracting text from: {document.file_url}") + extracted: ExtractedDocument = await download_and_extract( + file_url=document.file_url, + file_type=document.file_type, + ) + + if not extracted.full_text.strip(): + raise ValueError("No text content extracted from document") + + document.content = extracted.full_text + document.word_count = extracted.word_count + document.total_pages = extracted.total_pages + document.page_count = extracted.total_pages # keep backward compat + db.commit() + + # ── Step 3: Smart sentence-aware chunking with page tracking + logger.info(f"Chunking document ({extracted.total_pages} pages)...") + chunks: List[DocumentChunk] = _chunker.chunk_document( + full_text=extracted.full_text, + pages=extracted.pages, + ) + logger.info(f"Created {len(chunks)} chunks") + + # ── Step 4: Generate embeddings and persist ───────────────── + logger.info("Generating embeddings...") + for chunk in chunks: + embedding_values = await self._embed(chunk.text) + if embedding_values is None: + logger.warning( + f"Skipping chunk {chunk.chunk_index} — embedding failed" + ) + continue + + doc_embedding = DocumentEmbedding( + document_id=document.id, + chunk_index=chunk.chunk_index, + chunk_text=chunk.text, + embedding=embedding_values, # stored as vector(768) + page_numbers=chunk.page_numbers, # e.g. [12, 13] + section_title=chunk.section_title, # e.g. "Safety Procedures" + start_page=( + chunk.page_numbers[0] if chunk.page_numbers else None + ), + end_page=chunk.page_numbers[-1] if chunk.page_numbers else None, + embedding_model=self.embedding_model, + ) + db.add(doc_embedding) + + db.commit() + + # ── Step 5: Run multi-agent analysis via orchestrator ──────── + document.status = DocumentStatus.ANALYZING + db.commit() + + logger.info("Running AgentOrchestrator (4 agents)...") + results = await _orchestrator.analyze_document( + text=extracted.full_text, + pages=extracted.pages, + ) + + # ── Step 6: Map orchestrator results to document fields ────── + classification = results.get("classification", {}) + safety = results.get("safety", {}) + entities = results.get("entities", {}) + summary = results.get("summary", {}) + + # Classification + category_map = { + "safety_protocol": DocumentCategory.SAFETY_PROTOCOL, + "equipment_manual": DocumentCategory.EQUIPMENT_MANUAL, + "regulatory": DocumentCategory.REGULATORY, + "incident_report": DocumentCategory.INCIDENT_REPORT, + "geological": DocumentCategory.GEOLOGICAL, + "environmental": DocumentCategory.ENVIRONMENTAL, + "training": DocumentCategory.TRAINING, + "permit": DocumentCategory.PERMIT, + "maintenance": DocumentCategory.MAINTENANCE, + } + document.category = category_map.get( + classification.get("category"), DocumentCategory.OTHER + ) + document.subcategory = classification.get("subcategory") + document.classification_confidence = float( + classification.get("confidence", 0.5) + ) + + # Safety + status_map = { + "compliant": ComplianceStatus.COMPLIANT, + "warning": ComplianceStatus.WARNING, + "violation": ComplianceStatus.VIOLATION, + } + document.safety_score = ( + float(safety.get("score", 50)) + if safety.get("score") is not None + else None + ) + document.compliance_status = status_map.get( + safety.get("status"), ComplianceStatus.PENDING + ) + document.hazards_detected = safety.get("hazards", []) + document.safety_recommendations = safety.get("recommendations", []) + + # Entities & Summary + document.entities = entities if isinstance(entities, dict) else {} + document.summary = summary.get("summary", "Summary not available.") + document.key_points = summary.get("key_points", []) + + # ── Step 7: Mark completed ─────────────────────────────────── + document.status = DocumentStatus.COMPLETED + document.processed_at = datetime.now(timezone.utc) + db.commit() + + logger.info(f"Document processing completed: {document_id}") + return True + + except QuotaExceededError as qe: + # Quota hit: text extraction + embeddings already succeeded. + # Mark COMPLETED with partial data so re-analyze is available. + logger.error( + f"Gemini quota exceeded during agent analysis for {document_id}: {qe}" + ) + document.status = DocumentStatus.COMPLETED + document.processing_error = ( + "AI analysis incomplete: Gemini API quota exceeded. " + "Click Re-analyze to run the full analysis when quota resets." + ) + document.summary = "Summary not available — Gemini quota exceeded. Re-analyze to generate." + document.key_points = [] + document.processed_at = datetime.now(timezone.utc) + db.commit() + return False + + except Exception as e: + logger.error( + f"Document processing failed: {document_id} — {e}", exc_info=True + ) + document.status = DocumentStatus.FAILED + document.processing_error = str(e) + db.commit() + return False + + async def _embed(self, text: str) -> list | None: + """Generate a single embedding vector via Gemini embedding API.""" + try: + result = genai.embed_content( + model=self.embedding_model, + content=text, + task_type="retrieval_document", + output_dimensionality=768, + ) + return result["embedding"] + except Exception as e: + logger.warning(f"Embedding generation failed: {e}") + return None + + +# ── Background task wrapper ──────────────────────────────────────────────────── + + +async def process_document_async(document_id: str) -> None: + """Async wrapper used with FastAPI BackgroundTasks.""" + service = DocumentService() + await service.process_document(document_id) diff --git a/app/services/extractors.py b/app/services/extractors.py new file mode 100644 index 0000000000000000000000000000000000000000..16e738d1576fea530b7928bf6e6cf7ce817d5226 --- /dev/null +++ b/app/services/extractors.py @@ -0,0 +1,237 @@ +""" +Document Text Extractors +Page-aware text extraction for PDF, DOCX, and plain-text files. + +Replaces the monolithic _download_and_extract() in DocumentService with +dedicated extractor classes that track page boundaries — required for +context-aware answers that cite page numbers. +""" + +import io +import logging +import os +import tempfile +from dataclasses import dataclass, field +from typing import List, Optional + +import httpx + +logger = logging.getLogger(__name__) + + +# ── Data structures ──────────────────────────────────────────────────────────── + + +@dataclass +class PageContent: + """Text content of a single page.""" + + page_number: int # 1-indexed + text: str # raw text of this page + char_start: int = 0 # character offset where this page starts in full_text + char_end: int = 0 # character offset where this page ends + + +@dataclass +class ExtractedDocument: + """Result of text extraction from a document file.""" + + full_text: str + pages: List[PageContent] + total_pages: int + file_type: str + metadata: dict = field(default_factory=dict) # author, title, creation_date, etc. + + @property + def word_count(self) -> int: + return len(self.full_text.split()) + + +# ── Extractors ───────────────────────────────────────────────────────────────── + + +class PDFExtractor: + """ + Page-aware PDF text extraction using pypdf. + + Iterates page-by-page to build a per-page text map, which is then used + by ChunkingService to annotate each chunk with its page numbers. + """ + + def extract(self, file_path: str) -> ExtractedDocument: + from pypdf import PdfReader + + pages: List[PageContent] = [] + char_offset = 0 + + try: + reader = PdfReader(file_path) + for page_num, page in enumerate(reader.pages, start=1): + page_text = page.extract_text() or "" + pages.append( + PageContent( + page_number=page_num, + text=page_text, + char_start=char_offset, + char_end=char_offset + len(page_text), + ) + ) + char_offset += len(page_text) + + full_text = "".join(p.text for p in pages) + + # Extract PDF metadata + metadata = {} + if reader.metadata: + for key, val in reader.metadata.items(): + if val: + clean_key = key.lstrip("/").lower() + metadata[clean_key] = str(val) + + return ExtractedDocument( + full_text=full_text, + pages=pages, + total_pages=len(pages), + file_type="application/pdf", + metadata=metadata, + ) + + except Exception as e: + logger.error(f"pypdf extraction failed: {e}", exc_info=True) + raise ValueError(f"Failed to extract text from PDF: {e}") + + +class DocxExtractor: + """DOCX text extraction with paragraph tracking (approximate page numbers).""" + + def extract(self, file_path: str) -> ExtractedDocument: + from docx import Document as DocxDocument + + doc = DocxDocument(file_path) + paragraphs = [para.text for para in doc.paragraphs if para.text.strip()] + full_text = "\n".join(paragraphs) + + # DOCX doesn't expose true page numbers in python-docx. + # We approximate: every ~3000 chars ≈ 1 page. + pages = self._approximate_pages(full_text) + + return ExtractedDocument( + full_text=full_text, + pages=pages, + total_pages=len(pages), + file_type="application/vnd.openxmlformats-officedocument.wordprocessingml.document", + metadata={}, + ) + + def _approximate_pages( + self, text: str, chars_per_page: int = 3000 + ) -> List[PageContent]: + pages = [] + for i in range(0, max(1, len(text)), chars_per_page): + page_text = text[i : i + chars_per_page] + pages.append( + PageContent( + page_number=len(pages) + 1, + text=page_text, + char_start=i, + char_end=i + len(page_text), + ) + ) + return pages or [ + PageContent(page_number=1, text=text, char_start=0, char_end=len(text)) + ] + + +class PlainTextExtractor: + """Plain text extraction.""" + + def extract(self, file_path: str) -> ExtractedDocument: + with open(file_path, "r", encoding="utf-8", errors="ignore") as f: + full_text = f.read() + + # Approximate pages for plain text too + chars_per_page = 3000 + pages = [] + for i in range(0, max(1, len(full_text)), chars_per_page): + page_text = full_text[i : i + chars_per_page] + pages.append( + PageContent( + page_number=len(pages) + 1, + text=page_text, + char_start=i, + char_end=i + len(page_text), + ) + ) + + return ExtractedDocument( + full_text=full_text, + pages=pages + or [ + PageContent( + page_number=1, text=full_text, char_start=0, char_end=len(full_text) + ) + ], + total_pages=max(1, len(pages)), + file_type="text/plain", + ) + + +# ── File downloader + dispatcher ─────────────────────────────────────────────── + +EXTENSION_MAP = { + "application/pdf": ".pdf", + "application/vnd.openxmlformats-officedocument.wordprocessingml.document": ".docx", + "application/msword": ".doc", + "text/plain": ".txt", +} + +_extractors = { + "pdf": PDFExtractor(), + "docx": DocxExtractor(), + "txt": PlainTextExtractor(), +} + + +async def download_and_extract(file_url: str, file_type: str) -> ExtractedDocument: + """ + Download a file from URL and extract its text content with page tracking. + + Args: + file_url: Remote URL of the document file + file_type: MIME type string + + Returns: + ExtractedDocument with full text and per-page content + + Raises: + ValueError: If file type is not supported or extraction fails + """ + suffix = EXTENSION_MAP.get(file_type, ".tmp") + + # Validate URL scheme before downloading (SSRF prevention) + from urllib.parse import urlparse + + parsed = urlparse(file_url) + if parsed.scheme not in ("https", "http"): + raise ValueError(f"Unsupported URL scheme: {parsed.scheme}") + + async with httpx.AsyncClient(timeout=60.0, follow_redirects=True) as client: + response = await client.get(file_url) + response.raise_for_status() + + with tempfile.NamedTemporaryFile(delete=False, suffix=suffix) as tmp: + tmp.write(response.content) + tmp_path = tmp.name + + try: + if "pdf" in file_type.lower(): + return _extractors["pdf"].extract(tmp_path) + elif "docx" in file_type.lower() or "wordprocessingml" in file_type.lower(): + return _extractors["docx"].extract(tmp_path) + else: + return _extractors["txt"].extract(tmp_path) + finally: + try: + os.unlink(tmp_path) + except OSError: + pass diff --git a/app/services/hybrid_search.py b/app/services/hybrid_search.py new file mode 100644 index 0000000000000000000000000000000000000000..3c466202c78d0467de6d6c3bc6546463a0db8bc9 --- /dev/null +++ b/app/services/hybrid_search.py @@ -0,0 +1,267 @@ +""" +Hybrid Search Service + +Combines pgvector cosine similarity (semantic) with pg_trgm +trigram similarity (keyword/BM25) using Reciprocal Rank Fusion. + +This catches both: + - Semantic matches: "how to prevent methane explosions" + - Keyword matches: "30 CFR 75.323", "Caterpillar D11" + +Both search paths use the same PostgreSQL database — no external +services needed. +""" + +import logging +from collections import defaultdict +from typing import Any, Dict, List, Optional + +from sqlalchemy import text +from sqlalchemy.orm import Session + +from app.config import settings + +logger = logging.getLogger(__name__) + + +async def vector_search( + query_embedding: List[float], + db: Session, + user_id: str, + document_ids: Optional[List[str]] = None, + top_k: int = 20, +) -> List[Dict[str, Any]]: + """ + Pure pgvector cosine similarity search. + Returns chunks ranked by vector distance. + """ + doc_filter = "" + params: Dict[str, Any] = { + "user_id": user_id, + "embedding": query_embedding, + "top_k": top_k, + "threshold": settings.SIMILARITY_THRESHOLD, + } + + if document_ids: + doc_filter = "AND d.id = ANY(CAST(:doc_ids AS uuid[]))" + params["doc_ids"] = document_ids + + sql = text( + f""" + SELECT + de.id, + de.chunk_text, + de.page_numbers, + de.section_title, + de.start_page, + de.chunk_index, + d.id AS document_id, + d.title AS document_title, + d.file_name, + d.file_url, + 1 - (de.embedding <=> CAST(:embedding AS vector)) AS similarity + FROM document_embeddings de + JOIN documents d ON d.id = de.document_id + WHERE d.user_id = :user_id + AND d.status = 'COMPLETED' + {doc_filter} + AND (1 - (de.embedding <=> CAST(:embedding AS vector))) >= :threshold + ORDER BY de.embedding <=> CAST(:embedding AS vector) + LIMIT :top_k + """ + ) + + try: + rows = db.execute(sql, params).fetchall() + except Exception as e: + logger.error(f"Vector search failed: {e}", exc_info=True) + db.rollback() + return [] + + return [ + { + "id": str(row.id), + "text": row.chunk_text, + "page_numbers": row.page_numbers + or ([row.start_page] if row.start_page else []), + "section_title": row.section_title, + "chunk_index": row.chunk_index, + "document_id": str(row.document_id), + "document_title": row.document_title, + "file_name": row.file_name, + "file_url": row.file_url, + "score": float(row.similarity), + } + for row in rows + ] + + +def bm25_search( + query_text: str, + db: Session, + user_id: str, + document_ids: Optional[List[str]] = None, + top_k: int = 20, +) -> List[Dict[str, Any]]: + """ + PostgreSQL pg_trgm similarity search (keyword/BM25 equivalent). + Catches exact and partial string matches that vector search misses. + """ + doc_filter = "" + params: Dict[str, Any] = { + "user_id": user_id, + "query_text": query_text, + "top_k": top_k, + } + + if document_ids: + doc_filter = "AND d.id = ANY(CAST(:doc_ids AS uuid[]))" + params["doc_ids"] = document_ids + + sql = text( + f""" + SELECT + de.id, + de.chunk_text, + de.page_numbers, + de.section_title, + de.start_page, + de.chunk_index, + d.id AS document_id, + d.title AS document_title, + d.file_name, + d.file_url, + similarity(de.chunk_text, :query_text) AS bm25_score + FROM document_embeddings de + JOIN documents d ON d.id = de.document_id + WHERE d.user_id = :user_id + AND d.status = 'COMPLETED' + {doc_filter} + AND de.chunk_text % :query_text + ORDER BY bm25_score DESC + LIMIT :top_k + """ + ) + + try: + rows = db.execute(sql, params).fetchall() + except Exception as e: + # pg_trgm extension might not be available — fall back gracefully + logger.warning(f"BM25 search failed (pg_trgm may be unavailable): {e}") + db.rollback() + return [] + + return [ + { + "id": str(row.id), + "text": row.chunk_text, + "page_numbers": row.page_numbers + or ([row.start_page] if row.start_page else []), + "section_title": row.section_title, + "chunk_index": row.chunk_index, + "document_id": str(row.document_id), + "document_title": row.document_title, + "file_name": row.file_name, + "file_url": row.file_url, + "score": float(row.bm25_score), + } + for row in rows + ] + + +def reciprocal_rank_fusion( + list_a: List[Dict[str, Any]], + list_b: List[Dict[str, Any]], + k: int = None, +) -> List[Dict[str, Any]]: + """ + Reciprocal Rank Fusion (RRF) combines two ranked lists. + + RRF_score(d) = sum(1 / (k + rank_i(d))) for each list + + Args: + list_a: First ranked list (e.g., vector search results) + list_b: Second ranked list (e.g., BM25 results) + k: RRF constant (higher = less rank influence, default from config) + + Returns: + Merged list sorted by RRF score (descending) + """ + if k is None: + k = settings.RRF_K + + scores = defaultdict(float) + all_items = {} + + for rank, item in enumerate(list_a): + item_id = item["id"] + scores[item_id] += 1.0 / (k + rank + 1) + all_items[item_id] = item + + for rank, item in enumerate(list_b): + item_id = item["id"] + scores[item_id] += 1.0 / (k + rank + 1) + if item_id not in all_items: + all_items[item_id] = item + + ranked = sorted(scores.items(), key=lambda x: -x[1]) + return [all_items[item_id] for item_id, _ in ranked] + + +async def hybrid_search( + query_text: str, + query_embedding: List[float], + db: Session, + user_id: str, + document_ids: Optional[List[str]] = None, + top_k: int = None, +) -> List[Dict[str, Any]]: + """ + Hybrid search combining vector similarity + pg_trgm BM25 via RRF. + + Flow: + 1. Run vector search (pgvector cosine) → list A + 2. Run BM25 search (pg_trgm) → list B + 3. Combine via Reciprocal Rank Fusion + 4. Return top_k fused results + + If hybrid search is disabled or BM25 fails, falls back to vector-only. + """ + if top_k is None: + top_k = settings.RERANK_OVER_FETCH + + # Always run vector search + vector_results = await vector_search( + query_embedding=query_embedding, + db=db, + user_id=user_id, + document_ids=document_ids, + top_k=top_k, + ) + + if not settings.ENABLE_HYBRID_SEARCH: + return vector_results + + # Run BM25 search (may fail if pg_trgm unavailable) + bm25_results = bm25_search( + query_text=query_text, + db=db, + user_id=user_id, + document_ids=document_ids, + top_k=top_k, + ) + + if not bm25_results: + # BM25 returned nothing — fall back to vector-only + return vector_results + + # Fuse results + fused = reciprocal_rank_fusion(vector_results, bm25_results) + + logger.debug( + f"Hybrid search: {len(vector_results)} vector + " + f"{len(bm25_results)} BM25 → {len(fused)} fused" + ) + + return fused[:top_k] diff --git a/app/services/llm_provider.py b/app/services/llm_provider.py new file mode 100644 index 0000000000000000000000000000000000000000..c528532a96506c0e7c2e862376c061d64f8421cf --- /dev/null +++ b/app/services/llm_provider.py @@ -0,0 +1,39 @@ +""" +LLM Provider Initialization +Configures and exposes asynchronous clients for multiple AI providers. +""" + +import google.generativeai as genai +from cerebras.cloud.sdk import AsyncCerebras +from openai import AsyncOpenAI + +from app.config import settings + +# Initialize Groq client using OpenAI compatibility wrapper +groq_client = AsyncOpenAI( + base_url="https://api.groq.com/openai/v1", + api_key=settings.GROQ_API_KEY +) + +# Initialize Mistral client using OpenAI compatibility wrapper +mistral_client = AsyncOpenAI( + base_url="https://api.mistral.ai/v1", + api_key=settings.MISTRAL_API_KEY +) + +# Configure native Gemini API globally +genai.configure(api_key=settings.GEMINI_API_KEY) + +def get_groq_client() -> AsyncOpenAI: + return groq_client + +def get_mistral_client() -> AsyncOpenAI: + return mistral_client + +# Initialize Cerebras client +cerebras_client = AsyncCerebras( + api_key=settings.CEREBRAS_API_KEY +) + +def get_cerebras_client() -> AsyncCerebras: + return cerebras_client diff --git a/app/services/queue.py b/app/services/queue.py new file mode 100644 index 0000000000000000000000000000000000000000..fdcfcbdab762d02044391cca1a5098716a096cf0 --- /dev/null +++ b/app/services/queue.py @@ -0,0 +1,114 @@ +""" +Serverless Task Queue +Lightweight in-memory queue for document processing. +""" + +import asyncio +import logging +from typing import Optional + +logger = logging.getLogger(__name__) + +# Single global queue for the server process +_task_queue: Optional[asyncio.Queue] = None + + +def get_queue() -> asyncio.Queue: + global _task_queue + if _task_queue is None: + _task_queue = asyncio.Queue() + return _task_queue + + +async def document_worker(): + """Background worker that processes documents sequentially/concurrently.""" + from app.services.document_service import process_document_async + + queue = get_queue() + logger.info("Document worker started.") + + while True: + try: + document_id = await queue.get() + logger.info(f"Worker picked up document: {document_id}") + + try: + # We await the document processing. + # Concurrency is handled internally by Orchestrator or we can spawn tasks here. + # Since DocumentService is async, we can just await it directly or create a task. + asyncio.create_task(process_document_async(document_id)) + except Exception as e: + logger.error( + f"Worker failed dispatching document {document_id}: {e}", + exc_info=True, + ) + finally: + queue.task_done() + + except asyncio.CancelledError: + logger.info("Document worker cancelled.") + break + except Exception as e: + logger.error(f"Error in document worker loop: {e}", exc_info=True) + await asyncio.sleep(1) + + +def enqueue_document_task(document_id: str): + """Adds a document to the queue without blocking.""" + queue = get_queue() + try: + queue.put_nowait(document_id) + logger.info(f"Document {document_id} enqueued.") + except asyncio.QueueFull: + logger.error(f"Failed to enqueue document {document_id}: Queue is full") + + +# ── Compliance Audit Queue ───────────────────────────────────────────────────── + +_compliance_queue: Optional[asyncio.Queue] = None + + +def get_compliance_queue() -> asyncio.Queue: + global _compliance_queue + if _compliance_queue is None: + _compliance_queue = asyncio.Queue() + return _compliance_queue + + +async def compliance_worker(): + """Background worker that processes compliance audits.""" + from app.services.compliance_service import run_compliance_audit_async + + queue = get_compliance_queue() + logger.info("Compliance worker started.") + + while True: + try: + audit_id = await queue.get() + logger.info(f"Compliance worker picked up audit: {audit_id}") + + try: + asyncio.create_task(run_compliance_audit_async(audit_id)) + except Exception as e: + logger.error( + f"Worker failed dispatching audit {audit_id}: {e}", exc_info=True + ) + finally: + queue.task_done() + + except asyncio.CancelledError: + logger.info("Compliance worker cancelled.") + break + except Exception as e: + logger.error(f"Error in compliance worker loop: {e}", exc_info=True) + await asyncio.sleep(1) + + +async def enqueue_compliance_task(audit_id: str): + """Adds a compliance audit to the queue without blocking.""" + queue = get_compliance_queue() + try: + queue.put_nowait(audit_id) + logger.info(f"Compliance audit {audit_id} enqueued.") + except asyncio.QueueFull: + logger.error(f"Failed to enqueue audit {audit_id}: Queue is full") diff --git a/app/services/reranker.py b/app/services/reranker.py new file mode 100644 index 0000000000000000000000000000000000000000..5f516ad224804cb86d45694e7f3c46f746478542 --- /dev/null +++ b/app/services/reranker.py @@ -0,0 +1,102 @@ +""" +Cross-Encoder Reranking Service + +After initial vector/BM25 retrieval returns candidate chunks, +a cross-encoder model reads the query + each chunk together +and produces a true semantic relevance score. + +This catches cases where cosine similarity ranks a less-relevant +chunk higher than a more-relevant one. + +Model: cross-encoder/ms-marco-MiniLM-L-6-v2 + - Trained on MS MARCO passage ranking + - ~80M params, runs on CPU in ~50ms per batch of 20 + - Free, no API key needed +""" + +import logging +from typing import List + +from app.config import settings + +logger = logging.getLogger(__name__) + +_model = None + + +def _get_model(): + """Lazy-load the cross-encoder model (loaded once, cached globally).""" + global _model + if _model is None: + try: + from sentence_transformers import CrossEncoder + + logger.info(f"Loading reranker model: {settings.RERANK_MODEL}") + _model = CrossEncoder(settings.RERANK_MODEL) + logger.info("Reranker model loaded successfully") + except ImportError: + logger.error( + "sentence-transformers not installed. " + "Install with: pip install sentence-transformers" + ) + raise + return _model + + +def rerank( + query: str, + chunks: List[dict], + top_k: int = None, + text_key: str = "text", +) -> List[dict]: + """ + Rerank chunks by cross-encoder relevance score. + + Args: + query: The user's search query + chunks: List of chunk dicts, each must have `text_key` field + top_k: Number of top chunks to return (default: settings.RERANK_TOP_K) + text_key: Key in chunk dict containing the text to score + + Returns: + Top-k chunks sorted by cross-encoder score (descending). + Each chunk gets an added `rerank_score` field. + """ + if not chunks: + return [] + + if top_k is None: + top_k = settings.RERANK_TOP_K + + # If only 1 chunk, no reranking needed + if len(chunks) <= top_k: + for c in chunks: + c["rerank_score"] = c.get("score", 0.0) + return chunks + + model = _get_model() + + # Build query-document pairs for cross-encoder + pairs = [(query, chunk[text_key]) for chunk in chunks] + + try: + scores = model.predict(pairs) + except Exception as e: + logger.error(f"Reranking failed: {e}", exc_info=True) + # Fall back to original ordering + for c in chunks: + c["rerank_score"] = c.get("score", 0.0) + return chunks[:top_k] + + # Attach rerank scores and sort + for chunk, score in zip(chunks, scores): + chunk["rerank_score"] = float(score) + + chunks.sort(key=lambda c: c["rerank_score"], reverse=True) + + reranked = chunks[:top_k] + logger.debug( + f"Reranked {len(chunks)} chunks → top {top_k} " + f"(best score: {reranked[0]['rerank_score']:.4f})" + ) + return reranked diff --git a/app/workers/__init__.py b/app/workers/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..0938cb64e218e191f7259e844c85a03f5e80178b --- /dev/null +++ b/app/workers/__init__.py @@ -0,0 +1,4 @@ +""" +Background Workers +Celery task queue for async document processing +""" diff --git a/app/workers/celery_app.py b/app/workers/celery_app.py new file mode 100644 index 0000000000000000000000000000000000000000..5b247a30098fb14068baadb927af3c671dd0d5b1 --- /dev/null +++ b/app/workers/celery_app.py @@ -0,0 +1,41 @@ +""" +Celery Application +Background task queue for async document processing +""" + +import logging + +from celery import Celery + +from app.config import settings + +logger = logging.getLogger(__name__) + +# Create Celery app +celery_app = Celery( + "miningniti", + broker=settings.REDIS_URL, + backend=settings.REDIS_URL, + include=["app.workers.tasks"], +) + +# Celery configuration +celery_app.conf.update( + task_serializer="json", + accept_content=["json"], + result_serializer="json", + timezone="UTC", + enable_utc=True, + task_track_started=True, + task_acks_late=True, + worker_prefetch_multiplier=1, + task_routes={ + "app.workers.tasks.process_document_task": {"queue": "documents"}, + }, +) + + +@celery_app.task(name="app.workers.tasks.health_check") +def health_check(): + """Simple health check task""" + return {"status": "ok"} diff --git a/app/workers/tasks.py b/app/workers/tasks.py new file mode 100644 index 0000000000000000000000000000000000000000..9a9f688f1e58e7b368964c61ece75392be293fc5 --- /dev/null +++ b/app/workers/tasks.py @@ -0,0 +1,49 @@ +""" +Celery Tasks +Background processing tasks for document analysis +""" + +import asyncio +import logging + +from app.workers.celery_app import celery_app + +logger = logging.getLogger(__name__) + + +@celery_app.task( + bind=True, + name="app.workers.tasks.process_document_task", + max_retries=3, + default_retry_delay=30, +) +def process_document_task(self, document_id: str): + """ + Celery task for processing a document with AI agents. + Runs the async document processing pipeline in a sync context. + + Args: + document_id: UUID string of the document to process + """ + logger.info(f"Starting Celery task for document: {document_id}") + + try: + # Run the async processing pipeline in a new event loop + from app.services.document_service import DocumentService + + async def _run(): + service = DocumentService() + return await service.process_document(document_id) + + result = asyncio.run(_run()) + + if result: + logger.info(f"Document processed successfully: {document_id}") + else: + logger.error(f"Document processing returned False: {document_id}") + + return {"document_id": document_id, "success": result} + + except Exception as exc: + logger.error(f"Celery task failed for document {document_id}: {exc}") + raise self.retry(exc=exc) diff --git a/requirements.txt b/requirements.txt new file mode 100644 index 0000000000000000000000000000000000000000..6e7a822e1fc0f263cdb42db4d9e0f10bedc9fc07 --- /dev/null +++ b/requirements.txt @@ -0,0 +1,89 @@ +# ============================================================================= +# MiningNiti Backend — Requirements +# AI-Powered Document Intelligence for the Mining Industry +# ============================================================================= + +# ----------------------------------------------------------------------------- +# Core Framework +# ----------------------------------------------------------------------------- +fastapi==0.128.0 +uvicorn[standard]==0.40.0 +python-multipart==0.0.9 +python-dotenv==1.0.1 +pydantic==2.9.0 +pydantic-settings==2.5.0 + +# ----------------------------------------------------------------------------- +# Database & ORM +# ----------------------------------------------------------------------------- +sqlalchemy==2.0.35 +psycopg[binary,pool]==3.2.3 +psycopg2-binary==2.9.10 +pgvector==0.3.0 +alembic==1.13.0 + +# ----------------------------------------------------------------------------- +# AI/ML - Google Gemini, OpenAI compat +# ----------------------------------------------------------------------------- +google-generativeai==0.8.6 +google-auth==2.35.0 +google-api-python-client==2.140.0 +openai==1.50.0 +cerebras-cloud-sdk>=1.0.0 + +# ----------------------------------------------------------------------------- +# Vector Operations +# ----------------------------------------------------------------------------- +numpy==1.26.4 +scikit-learn==1.5.1 +sentence-transformers>=3.0.0 + +# ----------------------------------------------------------------------------- +# Document Processing +# ----------------------------------------------------------------------------- +pypdf==4.2.0 +python-docx==1.1.2 +python-magic==0.4.27 + +# ----------------------------------------------------------------------------- +# Background Jobs (Optional: used for production scaling, BackgroundTasks used in dev) +# ----------------------------------------------------------------------------- +# celery==5.6.2 # Uncomment for production Celery workers +redis==7.1.0 + +# ----------------------------------------------------------------------------- +# HTTP & Async +# ----------------------------------------------------------------------------- +httpx==0.28.1 +aiofiles==24.1.0 +requests==2.32.3 + +# ----------------------------------------------------------------------------- +# Security, Auth & Rate Limiting +# ----------------------------------------------------------------------------- +slowapi==0.1.9 +python-jose[cryptography]==3.3.0 +cryptography==46.0.3 +PyJWT==2.9.0 + +# ----------------------------------------------------------------------------- +# Utilities +# ----------------------------------------------------------------------------- +python-dateutil==2.9.0 +tqdm==4.66.4 + +# ----------------------------------------------------------------------------- +# Production Server +# ----------------------------------------------------------------------------- +gunicorn==23.0.0 +httptools>=0.6.3 + +# ----------------------------------------------------------------------------- +# Development & Testing +# ----------------------------------------------------------------------------- +pytest==8.3.0 +pytest-asyncio==0.24.0 +pytest-cov==5.0.0 +black==24.8.0 +isort==5.13.0 +mypy==1.11.0