added guardrails

Author: Trojan-Dev-AFK

README.md

@@ -41,20 +41,22 @@ Full setup guide: [docs/RUNBOOK.md](docs/RUNBOOK.md)
 | `POST` | `/v1/auth/logout` | Authenticated |
 | `GET` | `/v1/auth/me` | Authenticated |
 | `GET` | `/v1/auth/me/sessions` | Authenticated |
+| `GET` | `/v1/auth/sessions/company` | Admin (own company) or Super Admin |
 | `POST` | `/v1/companies/` | Super Admin |
 | `GET` | `/v1/companies/` | Super Admin |
 | `GET` | `/v1/companies/{id}` | Super Admin |
 | `PUT` | `/v1/companies/{id}` | Super Admin |
 | `DELETE` | `/v1/companies/{id}` | Super Admin |
-| `POST` | `/v1/users/` | Admin (own company) |
-| `GET` | `/v1/users/` | Admin (own company) |
-| `GET` | `/v1/users/{id}` | Admin (own company) |
-| `PUT` | `/v1/users/{id}` | Admin (own company) |
-| `DELETE` | `/v1/users/{id}` | Admin (own company) |
-| `POST` | `/v1/documents/upload` | Admin (own company) |
-| `GET` | `/v1/documents/` | Admin (own company) |
-| `GET` | `/v1/documents/{id}` | Admin (own company) |
-| `DELETE` | `/v1/documents/{id}` | Admin (own company) |
+| `POST` | `/v1/users/` | Admin (own company) or Super Admin |
+| `GET` | `/v1/users/` | Admin (own company) or Super Admin |
+| `GET` | `/v1/users/{id}` | Admin (own company) or Super Admin |
+| `PUT` | `/v1/users/{id}` | Admin (own company) or Super Admin |
+| `DELETE` | `/v1/users/{id}` | Admin (own company) or Super Admin |
+| `POST` | `/v1/documents/upload` | Admin (own company) or Super Admin (with company_id) |
+| `GET` | `/v1/documents/` | Admin (own company) or Super Admin |
+| `GET` | `/v1/documents/{id}` | Admin (own company) or Super Admin |
+| `DELETE` | `/v1/documents/{id}` | Admin (own company) or Super Admin |
+| `POST` | `/v1/chat/warmup` | Authenticated |
 | `POST` | `/v1/chat/invoke` | Admin or Employee |
 
 ---
@@ -101,6 +103,7 @@ ACCESS_TOKEN_EXPIRE_MINUTES=60
 EMBEDDING_MODEL=all-MiniLM-L6-v2
 CHUNK_SIZE=1000
 CHUNK_OVERLAP=200
+AGENT_WARMUP_ON_STARTUP=true
 
 # Encoding
 ENCODING=utf-8

app/agent/graph.py

@@ -46,6 +46,8 @@ def call_model(state: GraphState):
             " Do NOT just describe data in text.\n"
             "- First use 'search_documents' to gather the numbers if needed,"
             " then call 'generate_graph' with the data.\n"
+            "- Never call 'search_documents' repeatedly with the same query in a loop."
+            " If you already called it and have results, produce a final answer.\n"
             "- After calling 'generate_graph', keep your final text response brief (e.g. 'Here is your chart.'). "
             "Do NOT re-list all the data points in your text answer — the chart already shows them."
         )

app/agent/tools/vector_search.py

@@ -1,13 +1,15 @@
 """LangGraph tool: semantic search over document chunks using pgvector cosine distance."""
 
+from contextvars import ContextVar, Token
+
 from langchain_core.tools import tool
 from langchain_huggingface import HuggingFaceEmbeddings
 from sqlalchemy import select
 from sqlalchemy.exc import SQLAlchemyError
 
 from app.core.config import settings
 from app.core.logger import get_logger
-from app.db.models import DocumentChunk
+from app.db.models import Document, DocumentChunk
 from app.db.session import AsyncSessionLocal
 
 logger = get_logger(__name__)
@@ -17,6 +19,8 @@
 # every module. Eager loading here would load the model twice (once per process)
 # even though only the worker process ever runs queries.
 _embeddings_model: HuggingFaceEmbeddings | None = None
+_search_company_id: ContextVar[str | None] = ContextVar("search_company_id", default=None)
+_search_query_counts: ContextVar[dict[str, int] | None] = ContextVar("search_query_counts", default=None)
 
 
 def _get_embeddings_model() -> HuggingFaceEmbeddings:
@@ -29,13 +33,66 @@ def _get_embeddings_model() -> HuggingFaceEmbeddings:
     return _embeddings_model
 
 
+def warmup_vector_search() -> bool:
+    """Preload the embedding model and return True when this call performed the initial load."""
+    was_cold = _embeddings_model is None
+    _get_embeddings_model()
+    return was_cold
+
+
+def set_search_company_scope(company_id: str | None) -> Token:
+    """Set the tenant scope for vector search in the current request/task context."""
+    return _search_company_id.set(company_id)
+
+
+def reset_search_company_scope(token: Token) -> None:
+    """Reset the tenant scope for vector search in the current request/task context."""
+    _search_company_id.reset(token)
+
+
+def set_search_query_state() -> Token:
+    """Initialise per-request duplicate-query counters for vector search."""
+    return _search_query_counts.set({})
+
+
+def reset_search_query_state(token: Token) -> None:
+    """Reset per-request duplicate-query counters for vector search."""
+    _search_query_counts.reset(token)
+
+
 @tool
 async def search_documents(query: str) -> str:
     """
     Search the database for relevant document chunks based on a semantic query.
     Use this tool whenever you need to find factual information from the user's uploaded documents.
     """
-    logger.info("Vector search invoked", extra={"query_preview": query[:120]})
+    company_id = _search_company_id.get()
+    if not company_id:
+        logger.error("Vector search blocked — missing company scope")
+        return "Search context is unavailable for this request."
+
+    logger.info(
+        "Vector search invoked",
+        extra={"query_preview": query[:120], "company_id": company_id},
+    )
+
+    normalized_query = " ".join(query.lower().split())
+    query_counts = _search_query_counts.get()
+    if query_counts is None:
+        query_counts = {}
+        _search_query_counts.set(query_counts)
+
+    repeat_count = query_counts.get(normalized_query, 0)
+    if repeat_count >= 1:
+        logger.warning(
+            "Vector search blocked — duplicate query loop detected",
+            extra={"query_preview": query[:120], "company_id": company_id, "repeat_count": repeat_count},
+        )
+        return (
+            "You already ran this same search multiple times. "
+            "Do not call search_documents again for this query; provide your best final answer now."
+        )
+    query_counts[normalized_query] = repeat_count + 1
 
     try:
         query_vector = _get_embeddings_model().embed_query(query)
@@ -47,6 +104,8 @@ async def search_documents(query: str) -> str:
         async with AsyncSessionLocal() as db:
             stmt = (
                 select(DocumentChunk.text_content)
+                .join(Document, Document.id == DocumentChunk.document_id)
+                .where(Document.company_id == company_id)
                 .order_by(DocumentChunk.embedding.cosine_distance(query_vector))
                 .limit(5)
             )
@@ -60,5 +119,5 @@ async def search_documents(query: str) -> str:
         logger.info("Vector search returned no results", extra={"query_preview": query[:120]})
         return "No relevant information found in the documents."
 
-    logger.info("Vector search completed", extra={"results": len(chunks)})
+    logger.info("Vector search completed", extra={"results": len(chunks), "company_id": company_id})
     return "\n\n---\n\n".join(chunks)

app/api/v1/endpoints/auth.py

@@ -3,17 +3,17 @@
 import uuid
 from datetime import UTC, datetime, timedelta
 
-from fastapi import APIRouter, Depends, HTTPException, Request, status
+from fastapi import APIRouter, Depends, HTTPException, Query, Request, status
 from fastapi.security import OAuth2PasswordRequestForm
 from jose import JWTError, jwt
 from sqlalchemy import select
 from sqlalchemy.ext.asyncio import AsyncSession
 
-from app.api.dependencies import get_current_user, oauth2_scheme
+from app.api.dependencies import get_current_user, oauth2_scheme, require_admin_or_super_admin
 from app.core.config import settings
 from app.core.logger import get_logger
 from app.core.security import create_access_token, verify_password
-from app.db.models import TokenSession, User
+from app.db.models import TokenSession, User, UserRole
 from app.db.session import get_db
 from app.schemas.auth import LogoutResponse, MeResponse, TokenResponse
 from app.schemas.sessions import SessionResponse
@@ -157,3 +157,56 @@ async def list_my_sessions(
         extra={"user_id": current_user.id, "session_count": len(sessions)},
     )
     return sessions
+
+
+@router.get("/sessions/company", response_model=list[SessionResponse])
+async def list_company_sessions(
+    company_id: str | None = Query(
+        default=None,
+        description="Filter by company UUID. Admin users are always restricted to their own company.",
+    ),
+    current_user: User = Depends(require_admin_or_super_admin),
+    db: AsyncSession = Depends(get_db),
+):
+    """
+    List token sessions for a company.
+
+    **admin**: sees sessions for their own company only.
+    **super_admin**: can list all company sessions or filter by ``company_id``.
+    """
+    target_company_id = company_id
+    if current_user.role == UserRole.ADMIN:
+        if company_id and company_id != current_user.company_id:
+            logger.warning(
+                "Company sessions access denied",
+                extra={
+                    "user_id": current_user.id,
+                    "actor_company": current_user.company_id,
+                    "requested_company": company_id,
+                },
+            )
+            raise HTTPException(status_code=status.HTTP_403_FORBIDDEN, detail="Access denied")
+        target_company_id = current_user.company_id
+
+    stmt = select(TokenSession).join(User, User.id == TokenSession.user_id)
+
+    if target_company_id:
+        stmt = stmt.where(User.company_id == target_company_id)
+    else:
+        # "All company sessions" intentionally excludes super_admin accounts (company_id is NULL).
+        stmt = stmt.where(User.company_id.is_not(None))
+
+    stmt = stmt.order_by(TokenSession.issued_at.desc())
+    result = await db.execute(stmt)
+    sessions = result.scalars().all()
+
+    logger.info(
+        "Company sessions listed",
+        extra={
+            "actor": current_user.id,
+            "actor_role": str(current_user.role),
+            "target_company_id": target_company_id or "all_companies",
+            "session_count": len(sessions),
+        },
+    )
+    return sessions

app/api/v1/endpoints/chat.py

@@ -1,22 +1,68 @@
 """Chat endpoint — submits a query to the LangGraph RAG agent."""
 
+import asyncio
 import json
 import time
 
-from fastapi import APIRouter, Depends
+from fastapi import APIRouter, Depends, HTTPException, status
 from langchain_core.messages import HumanMessage, ToolMessage
+from langgraph.errors import GraphRecursionError
 
 from app.agent.graph import app_graph
-from app.api.dependencies import require_company_user
+from app.agent.tools.vector_search import (
+    reset_search_company_scope,
+    reset_search_query_state,
+    set_search_company_scope,
+    set_search_query_state,
+    warmup_vector_search,
+)
+from app.api.dependencies import get_current_user, require_company_user
 from app.core.exceptions import AgentError
 from app.core.logger import get_logger
 from app.db.models import User
-from app.schemas.chat import ChatRequest, ChatResponse
+from app.schemas.chat import ChatRequest, ChatResponse, WarmupResponse
 
 logger = get_logger(__name__)
 router = APIRouter()
 
 
+@router.post(
+    "/warmup",
+    response_model=WarmupResponse,
+    responses={
+        502: {"description": "Warmup failed — embedding model could not be loaded."},
+    },
+)
+async def warmup_agent(current_user: User = Depends(get_current_user)):
+    """Warm up agent resources so the next chat request avoids embedding-model cold start latency."""
+    start = time.monotonic()
+    try:
+        loaded_now = await asyncio.to_thread(warmup_vector_search)
+    except Exception as exc:
+        elapsed = time.monotonic() - start
+        logger.error(
+            "Agent warmup failed",
+            extra={"user_id": current_user.id, "elapsed_s": round(elapsed, 3)},
+            exc_info=exc,
+        )
+        raise AgentError("Agent warmup failed. Please try again.") from exc
+
+    elapsed = time.monotonic() - start
+    logger.info(
+        "Agent warmup completed",
+        extra={
+            "user_id": current_user.id,
+            "elapsed_s": round(elapsed, 3),
+            "embeddings_loaded_now": loaded_now,
+        },
+    )
+    return WarmupResponse(
+        message="Agent warmup completed",
+        embeddings_loaded_now=loaded_now,
+        elapsed_seconds=round(elapsed, 3),
+    )
+
+
 def _extract_graph_payload(messages) -> dict | None:
     for msg in reversed(messages):
         if isinstance(msg, ToolMessage) and msg.name == "generate_graph":
@@ -61,10 +107,36 @@ async def invoke_agent(
         },
     )
 
+    if not current_user.company_id:
+        logger.warning(
+            "Chat blocked — user has no company scope",
+            extra={"user_id": current_user.id, "role": str(current_user.role)},
+        )
+        raise HTTPException(
+            status_code=status.HTTP_403_FORBIDDEN,
+            detail="You do not have permission to perform this action.",
+        )
+
     start = time.monotonic()
+    search_scope_token = set_search_company_scope(current_user.company_id)
+    search_query_state_token = set_search_query_state()
     try:
         initial_state = {"messages": [HumanMessage(content=request.query)]}
-        final_state = await app_graph.ainvoke(initial_state)
+        final_state = await app_graph.ainvoke(initial_state, config={"recursion_limit": 8})
+    except GraphRecursionError as exc:
+        elapsed = time.monotonic() - start
+        logger.error(
+            "Agent halted by recursion limit",
+            extra={
+                "user_id": current_user.id,
+                "company_id": current_user.company_id,
+                "elapsed_s": round(elapsed, 3),
+            },
+            exc_info=exc,
+        )
+        raise AgentError(
+            "The agent could not converge on an answer. Please rephrase your question with more specific details."
+        ) from exc
     except Exception as exc:
         elapsed = time.monotonic() - start
         logger.error(
@@ -77,6 +149,9 @@ async def invoke_agent(
             exc_info=exc,
         )
         raise AgentError("The agent failed to process your request. Please try again.") from exc
+    finally:
+        reset_search_query_state(search_query_state_token)
+        reset_search_company_scope(search_scope_token)
 
     elapsed = time.monotonic() - start
     final_message = final_state["messages"][-1].content