From 32b08f8b0c227b7b53c268cbf9deba3b7fbf8d07 Mon Sep 17 00:00:00 2001 From: nickviljoen Date: Mon, 30 Mar 2026 18:29:35 +0200 Subject: [PATCH] Add Prompt Audit & Auto-Classification feature with Gemini integration Adds Gemini-powered agent classification system that analyzes agent instructions to determine category, risk level, discipline, and client detection. Includes admin Prompt Audit tab, audit review workflow, and auto-classification on sync. Co-Authored-By: Claude Opus 4.6 (1M context) --- .gitignore | 5 + CLAUDE.md | 54 +++- audit_analyzer.py | 461 ++++++++++++++++++++++++++++++ database.py | 3 + main.py | 81 +++++- models.py | 12 + requirements.txt | 1 + templates/admin/dashboard.html | 508 ++++++++++++++++++++++++++++++++- templates/agent_register.html | 52 ++-- 9 files changed, 1141 insertions(+), 36 deletions(-) create mode 100644 audit_analyzer.py diff --git a/.gitignore b/.gitignore index 43e1485..fc66aa5 100644 --- a/.gitignore +++ b/.gitignore @@ -97,3 +97,8 @@ Desktop.ini # MongoDB dumps *.dump *.bson + +# Planning/debug files +PLAN-*.md +CHANGELOG-*.md +check_db.py diff --git a/CLAUDE.md b/CLAUDE.md index c6a753a..e7adba9 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -27,6 +27,11 @@ Create `.env` file with: - `ALGORITHM`: JWT algorithm (default: HS256) - `ACCESS_TOKEN_EXPIRE_MINUTES`: Token expiration (default: 60) +#### Optional: Prompt Audit (Gemini) +- `GOOGLE_API_KEY`: Google API key for Gemini (if not set, audit is silently disabled) +- `AUDIT_GEMINI_MODEL`: Gemini model name (default: `gemini-2.5-pro`) +- `AUDIT_CONCURRENCY`: Batch size for sequential processing (default: 2) + #### Optional: Email Notifications (Mailgun) - `MAILGUN_API_KEY`: Mailgun API key (if not set, notifications are silently disabled) - `MAILGUN_DOMAIN`: Mailgun sending domain (e.g. your-domain.mailgun.org) @@ -68,6 +73,7 @@ Create `.env` file with: - `UserUpdate`: Includes `role` field for three-tier role management - `AiAgentCreate/AiAgentResponse`: API request/response models (includes `total_tokens`, `prompt_tokens`, `completion_tokens`, `discipline`, `rating`, `rating_count`, `client`, `client_name`, `studio_name`, `verification_status`, `verified_by`, `verified_date`) - `AgentCollectorCreate`: Collector API input model (includes `total_tokens`, `prompt_tokens`, `completion_tokens`, `discipline`, `client`, `client_name`, `studio_name`) +- `AuditReviewRequest`: Audit review input (`audit_status`: flagged/reviewed/cleared, `reviewer_notes`) - `AgentUsageStatsResponse`: Usage statistics response (includes `total_tokens`, `prompt_tokens`, `completion_tokens`) **crud.py**: Database operations using Motor (async MongoDB driver): @@ -77,8 +83,19 @@ Create `.env` file with: - All operations use ObjectId for MongoDB document IDs **database.py**: MongoDB connection setup with Motor async client -- Collections: `users`, `agents`, `agent_usage`, `token_notifications`, `agent_ratings` -- `ensure_indexes()`: Creates compound unique index on `agent_ratings(agent_id, user_id)` and index on `verification_status` +- Collections: `users`, `agents`, `agent_usage`, `token_notifications`, `agent_ratings`, `audit_history` +- `ensure_indexes()`: Creates compound unique index on `agent_ratings(agent_id, user_id)`, indexes on `verification_status`, `audit_status`, and `audit_history(agent_id, audit_date)` + +**audit_analyzer.py**: Gemini-powered agent classification and audit system: +- `is_gemini_configured()`: Returns False if `GOOGLE_API_KEY` not set (gracefully disabled) +- `analyze_single_agent()`: Sends agent instructions to Gemini 2.5 Pro, returns structured JSON with category, discipline, department, client detection, risk level, flags, and recommendations. Includes retry with exponential backoff for rate limits. +- `store_audit_result()`: Writes audit results to agent document (`audit_status`, `audit_category`, `audit_risk_level`, etc.) and inserts into `audit_history` collection +- `apply_classification_fields()`: Auto-assigns `discipline` (from defined list), `agent_department` (free text inferred from instructions), and client detection (`client = "yes"`, `verification_status = "needs_verification"`) — only overwrites fields that are currently empty/null +- `classify_single_agent()`: Convenience function for post-sync automatic classification. Loads agent, analyses, stores result, applies fields. +- `run_audit_batch()`: Batch processes agents sequentially with rate-limit-safe pauses. Supports `unclassified_only` and `single_agent_id` params. Returns summary with audited/failed/skipped counts. +- `get_all_audit_results()`: Returns all agents with audit fields for the Prompt Audit tab +- `update_audit_review()`: Admin marks audit as reviewed/cleared with notes +- Uses `google-genai` SDK (new API), model configurable via `AUDIT_GEMINI_MODEL` env var (default: `gemini-2.5-pro`) **notifications.py**: Mailgun email notification system: - `is_mailgun_configured()`: Returns False if env vars not set (gracefully disabled) @@ -200,6 +217,29 @@ Located in `templates/` directory: - Cooldown tracking in MongoDB `token_notifications` collection (default 24h, configurable) - Sends to all active admin users' email addresses +### Prompt Audit & Auto-Classification (Gemini) +- Automated analysis of agent `instructions` (system prompts) using Google Gemini 2.5 Pro +- **Two trigger modes:** + 1. **Automatic post-sync**: After the collector API (`POST /agents`) creates/updates an agent with instructions, a background task (`asyncio.create_task`) auto-classifies it. Non-blocking — sync response is not delayed. + 2. **Manual batch**: Admin clicks "Run Full Audit" or "Run Unclassified Only" on the Prompt Audit tab. Processes agents sequentially with 4-second pauses between batches to avoid Gemini rate limits. +- **Classification outputs per agent:** + - `audit_category`: Cat 1 (Internal Sandbox), Cat 1B (High Cost Internal), Cat 2 (Client-Exposed), Cat 3 (Client-Sold) + - `audit_risk_level`: low / medium / high / critical + - `audit_discipline`: Picks from existing discipline list (Strategy, Creative, Oversight including delivery, Optimization, Back Office including operations, Pencil Agents) + - `audit_department`: Free text inferred from instructions (e.g., "Project Management", "Media") + - `audit_is_client_work`: Boolean — detects client names, brands, client deliverables in instructions + - `audit_flags`: Array of risk flags (client_facing, handles_pii, uses_external_tools, etc.) + - `audit_summary`, `audit_recommendations`, `audit_category_reasoning`, `audit_discipline_reasoning`, `audit_client_work_reasoning`, `audit_client_name_detected` +- **Auto-assignment**: Gemini results auto-populate `discipline` and `agent_department` fields on the agent document (only if currently empty, to respect manual edits) +- **Client work auto-detection**: When `audit_is_client_work = true`, auto-sets `client = "yes"` and `verification_status = "needs_verification"` — agent appears on Verification tab. Does NOT overwrite manually-set values. +- **Audit statuses**: All audited agents start as `flagged`. Admin can mark as `reviewed` or `cleared` via the detail modal with optional notes. +- **Prompt Audit tab** on admin dashboard: Summary cards (Audited, Flagged, Reviewed, Cleared, Client Detected, No Instructions), filterable results table, detail modal with full analysis and review controls +- **Rate limit handling**: Retry with exponential backoff (10s, 20s, 40s) for 429/quota errors +- **Logging**: Uses Python `logging` module (`audit_analyzer` logger) for systemd journal visibility +- Agents without `instructions` are skipped (counted as "No Instructions") +- `audit_history` collection stores historical record of each audit run and review action +- Gracefully disabled when `GOOGLE_API_KEY` not set — UI shows config warning, buttons hidden + ### User Management - User registration with validation - Admin user creation capabilities @@ -214,7 +254,7 @@ Located in `templates/` directory: - Async operations using Motor driver - Indexed queries for performance - Data aggregation for statistics -- Collections: `users`, `agents`, `agent_usage`, `token_notifications`, `agent_ratings` +- Collections: `users`, `agents`, `agent_usage`, `token_notifications`, `agent_ratings`, `audit_history` ## Development Guidelines @@ -262,6 +302,7 @@ Key dependencies from requirements.txt: - **python-multipart**: Form handling - **requests**: HTTP client (used for Mailgun API calls) - **apscheduler**: Task scheduling (weekly digest email) +- **google-genai**: Google Gemini API client (used for prompt audit auto-classification) ## API Endpoints (New) @@ -270,4 +311,9 @@ Key dependencies from requirements.txt: - `PUT /api/admin/agents/{agent_id}/verify` — Approve/verify an agent (admin only) ### Weekly Digest -- `POST /api/admin/digest/send` — Manually trigger the weekly agent digest email (admin only) \ No newline at end of file +- `POST /api/admin/digest/send` — Manually trigger the weekly agent digest email (admin only) + +### Prompt Audit +- `POST /api/admin/audit/run` — Run Gemini audit batch. Optional JSON body: `{agent_id, unclassified_only}`. Returns `{status, total, audited_count, failed_count, skipped_count, results_summary}` (admin only) +- `GET /api/admin/audit/results` — Get all agents with audit data + `config_status.gemini_configured` (admin + readonly_admin) +- `PUT /api/admin/audit/{agent_id}/review` — Mark audit as reviewed/cleared with `{audit_status, reviewer_notes}` (admin only) \ No newline at end of file diff --git a/audit_analyzer.py b/audit_analyzer.py new file mode 100644 index 0000000..8de2b28 --- /dev/null +++ b/audit_analyzer.py @@ -0,0 +1,461 @@ +import os +import json +import re +import asyncio +import logging +from datetime import datetime +from bson import ObjectId + +from google import genai + +from database import agents_collection, audit_history_collection + +logger = logging.getLogger("audit_analyzer") +logger.setLevel(logging.INFO) +if not logger.handlers: + handler = logging.StreamHandler() + handler.setFormatter(logging.Formatter("%(asctime)s [AUDIT] %(message)s")) + logger.addHandler(handler) + + +def is_gemini_configured() -> bool: + return bool(os.getenv("GOOGLE_API_KEY")) + + +def _get_client(): + api_key = os.getenv("GOOGLE_API_KEY") + return genai.Client(api_key=api_key) + + +def _get_model_name(): + return os.getenv("AUDIT_GEMINI_MODEL", "gemini-2.5-pro") + + +SYSTEM_INSTRUCTION = """You are an AI agent compliance analyst and classifier. Analyse AI agents based on their +system prompts/instructions and metadata. You must: +1. Classify into business risk categories +2. Assign a business discipline +3. Infer the department/team from the instructions +4. Detect whether the agent is used for client-specific work + +Respond ONLY with valid JSON matching this schema: +{ + "category": "1" | "1B" | "2" | "3", + "category_reasoning": "string", + "discipline": "Strategy" | "Creative" | "Oversight including delivery" | "Optimization" | "Back Office including operations" | "Pencil Agents", + "discipline_reasoning": "string - why this discipline was chosen", + "department": "string or null - inferred team/department from instructions (e.g. Project Management, Creative Services, Finance, Media, Strategy). null if not determinable", + "is_client_work": true | false, + "client_work_reasoning": "string - evidence from instructions that this agent handles client-specific work, references client names, brands, deliverables, or external-facing outputs. Empty string if not client work", + "client_name_detected": "string or null - specific client or brand name found in instructions, null if none", + "flags": ["array", "of", "strings"], + "summary": "2-3 sentence analysis", + "recommendations": "which team(s) should review and why", + "risk_level": "low" | "medium" | "high" | "critical" +} + +CATEGORY DEFINITIONS: +Cat 1 - Internal Sandbox/Experimentation (Oliver AI Sandbox, behind the scenes, needs IT/Compliance) +Cat 1B - High Cost Internal (may incur large cost, not Cat 2 or 3) +Cat 2 - Client-Exposed Not Sold (Pencil platform, exposed to clients, needs Legal) +Cat 3 - Client-Sold (Pencil platform, sold to clients, needs Commercial team) + +DISCIPLINE DEFINITIONS (pick the best fit): +- Strategy: Agents focused on strategic planning, research, market analysis, insights +- Creative: Agents focused on creative work, content creation, design, copywriting +- Oversight including delivery: Agents focused on project management, delivery, QA, oversight, compliance +- Optimization: Agents focused on performance optimization, data analysis, efficiency, media planning +- Back Office including operations: Agents focused on internal operations, HR, finance, IT support, admin tasks +- Pencil Agents: Agents built on or for the Pencil platform specifically + +DEPARTMENT INFERENCE: +Look for clues in the instructions about which team or role uses this agent. Examples: +- "I am a project manager" -> department: "Project Management" +- "help the media team" -> department: "Media" +- "creative brief" -> department: "Creative" +- If no department/team clues found, set department to null + +CLIENT WORK DETECTION: +An agent is client work (is_client_work = true) if the instructions reference: +- Specific client names or brand names +- Client deliverables, client presentations, client reports +- External-facing outputs meant for clients +- Client briefs, client feedback, client approvals +- Work explicitly described as "for the client" or "client-facing" +Do NOT flag as client work if the agent merely mentions "users" or "stakeholders" generically. + +RISK LEVELS: +- low: Internal-only, limited capabilities +- medium: Internal with external tool access or moderate cost +- high: Client-facing or accesses sensitive data +- critical: Client-sold or handles financial/legal/PII data + +FLAGS TO CONSIDER: +internal_only, experimental, sandbox, client_facing, pencil_platform, +revenue_generating, not_for_sale, uses_external_tools, uses_code_interpreter, +uses_file_search, accesses_sensitive_data, handles_pii, high_cost, +resource_intensive, legal_review_needed, commercial_review_needed, +compliance_review_needed, no_instructions""" + + +USER_PROMPT_TEMPLATE = """Analyse this AI agent: + +AGENT NAME: {name} +DESCRIPTION: {description} +AUTHOR: {author} + +INSTRUCTIONS (SYSTEM PROMPT): +--- +{instructions} +--- + +TOOLS: {tools}""" + + +def _parse_json_response(text: str) -> dict: + """Parse JSON from Gemini response, with regex fallback.""" + # Try direct parse first + try: + return json.loads(text) + except json.JSONDecodeError: + pass + + # Regex fallback: extract JSON block + match = re.search(r'\{[\s\S]*\}', text) + if match: + try: + return json.loads(match.group()) + except json.JSONDecodeError: + pass + + # Last resort: return raw text as summary + return { + "category": "1", + "category_reasoning": "Could not parse LLM response", + "discipline": "Back Office including operations", + "discipline_reasoning": "Default — LLM response unparseable", + "department": None, + "is_client_work": False, + "client_work_reasoning": "", + "client_name_detected": None, + "flags": ["parse_error"], + "summary": text[:500], + "recommendations": "Manual review required — automated analysis failed to produce structured output", + "risk_level": "medium" + } + + +async def analyze_single_agent(agent_name: str, instructions: str, tools: str, + description: str, author: str, max_retries: int = 3) -> dict: + """Analyse a single agent using Gemini and return structured results. + Includes retry with exponential backoff for rate limits.""" + client = _get_client() + model_name = _get_model_name() + + prompt = USER_PROMPT_TEMPLATE.format( + name=agent_name, + description=description or "No description provided", + author=author or "Unknown", + instructions=instructions or "No instructions available", + tools=tools or "None specified" + ) + + for attempt in range(max_retries): + try: + response = await asyncio.to_thread( + client.models.generate_content, + model=model_name, + contents=prompt, + config=genai.types.GenerateContentConfig( + system_instruction=SYSTEM_INSTRUCTION + ) + ) + result = _parse_json_response(response.text) + result["agent_name"] = agent_name + return result + except Exception as e: + error_str = str(e).lower() + is_rate_limit = "429" in error_str or "rate" in error_str or "quota" in error_str or "resource_exhausted" in error_str + if is_rate_limit and attempt < max_retries - 1: + wait_time = (2 ** attempt) * 10 # 10s, 20s, 40s + logger.warning(f"Rate limited on '{agent_name}', retrying in {wait_time}s (attempt {attempt + 1}/{max_retries})") + await asyncio.sleep(wait_time) + continue + # Final attempt failed or non-rate-limit error + logger.error(f"Failed to analyse '{agent_name}': {e}") + e_final = e + + return { + "error": str(e_final), + "agent_name": agent_name, + "category": "1", + "category_reasoning": f"Analysis failed: {e_final}", + "discipline": None, + "discipline_reasoning": f"Analysis failed: {e_final}", + "department": None, + "is_client_work": False, + "client_work_reasoning": "", + "client_name_detected": None, + "flags": ["analysis_error"], + "summary": f"Automated analysis failed: {e_final}", + "recommendations": "Manual review required", + "risk_level": "medium" + } + + +async def store_audit_result(agent_id: str, audit_data: dict): + """Store audit results on the agent document and in audit_history.""" + now = datetime.utcnow() + + # Update agent document with audit fields + update_fields = { + "audit_status": "flagged", + "audit_date": now.isoformat(), + "audit_category": audit_data.get("category"), + "audit_risk_level": audit_data.get("risk_level"), + "audit_summary": audit_data.get("summary"), + "audit_flags": audit_data.get("flags", []), + "audit_recommendations": audit_data.get("recommendations"), + "audit_category_reasoning": audit_data.get("category_reasoning"), + "audit_discipline": audit_data.get("discipline"), + "audit_discipline_reasoning": audit_data.get("discipline_reasoning"), + "audit_department": audit_data.get("department"), + "audit_is_client_work": audit_data.get("is_client_work", False), + "audit_client_work_reasoning": audit_data.get("client_work_reasoning", ""), + "audit_client_name_detected": audit_data.get("client_name_detected"), + } + + await agents_collection.update_one( + {"_id": ObjectId(agent_id)}, + {"$set": update_fields} + ) + + # Insert into audit_history for historical tracking + history_entry = { + "agent_id": agent_id, + "audit_date": now, + **audit_data + } + await audit_history_collection.insert_one(history_entry) + + +async def apply_classification_fields(agent_id: str, audit_data: dict): + """Apply discipline, department, and client detection to the agent document. + + Only overwrites fields that are currently empty/null, to respect manual edits. + """ + agent = await agents_collection.find_one({"_id": ObjectId(agent_id)}) + if not agent: + return + + update_fields = {} + + # Auto-assign discipline if not already set + discipline = audit_data.get("discipline") + if discipline and not agent.get("discipline"): + update_fields["discipline"] = discipline + + # Auto-assign department if not already set + department = audit_data.get("department") + if department and not agent.get("agent_department"): + update_fields["agent_department"] = department + + # Auto-flag client work if not already manually set + is_client_work = audit_data.get("is_client_work", False) + if is_client_work and not agent.get("client"): + update_fields["client"] = "yes" + update_fields["verification_status"] = "needs_verification" + detected_name = audit_data.get("client_name_detected") + if detected_name and not agent.get("client_name"): + update_fields["client_name"] = detected_name + + if update_fields: + await agents_collection.update_one( + {"_id": ObjectId(agent_id)}, + {"$set": update_fields} + ) + + +async def classify_single_agent(agent_id: str): + """Convenience function for post-sync: load agent, analyse, store, apply fields.""" + if not is_gemini_configured(): + return + + try: + agent = await agents_collection.find_one({"_id": ObjectId(agent_id)}) + if not agent: + logger.warning(f"Agent {agent_id} not found") + return + + instructions = agent.get("instructions") + if not instructions: + logger.info(f"Agent '{agent.get('agent_name')}' has no instructions, skipping") + return + + # Build tools string from agent metadata + tools = ", ".join(agent.get("agent_capabilities", []) or []) + + audit_data = await analyze_single_agent( + agent_name=agent.get("agent_name", "Unknown"), + instructions=instructions, + tools=tools, + description=agent.get("agent_description", ""), + author=agent.get("created_by", "") + ) + + await store_audit_result(agent_id, audit_data) + await apply_classification_fields(agent_id, audit_data) + + logger.info(f"Classified '{agent.get('agent_name')}' as Cat {audit_data.get('category')}, " + f"discipline={audit_data.get('discipline')}, client={audit_data.get('is_client_work')}") + + except Exception as e: + logger.error(f"Failed to classify agent {agent_id}: {e}") + + +async def run_audit_batch(unclassified_only: bool = False, single_agent_id: str = None, + concurrency: int = None) -> dict: + """Run audit on multiple agents with concurrency control. + + Returns summary dict with counts and results. + """ + if not is_gemini_configured(): + return {"error": "GOOGLE_API_KEY not configured"} + + if concurrency is None: + concurrency = int(os.getenv("AUDIT_CONCURRENCY", "2")) + + # Build query + query = {} + if single_agent_id: + query["_id"] = ObjectId(single_agent_id) + elif unclassified_only: + query["audit_status"] = {"$exists": False} + + agents = await agents_collection.find(query).to_list(length=None) + + total = len(agents) + audited = 0 + failed = 0 + skipped = 0 + results = [] + + # Filter out agents without instructions first + agents_with_instructions = [] + for agent in agents: + if not agent.get("instructions"): + skipped += 1 + else: + agents_with_instructions.append(agent) + + logger.info(f"Audit batch: {total} total, {len(agents_with_instructions)} with instructions, {skipped} skipped") + + # Process agents sequentially in small batches to respect Gemini rate limits + batch_size = concurrency + for i in range(0, len(agents_with_instructions), batch_size): + batch = agents_with_instructions[i:i + batch_size] + + for agent in batch: + agent_id = str(agent["_id"]) + tools = ", ".join(agent.get("agent_capabilities", []) or []) + + try: + audit_data = await analyze_single_agent( + agent_name=agent.get("agent_name", "Unknown"), + instructions=agent.get("instructions"), + tools=tools, + description=agent.get("agent_description", ""), + author=agent.get("created_by", "") + ) + + if audit_data.get("error"): + failed += 1 + results.append({"agent_name": agent.get("agent_name"), "error": audit_data["error"]}) + else: + await store_audit_result(agent_id, audit_data) + await apply_classification_fields(agent_id, audit_data) + audited += 1 + results.append({ + "agent_name": agent.get("agent_name"), + "category": audit_data.get("category"), + "discipline": audit_data.get("discipline"), + "risk_level": audit_data.get("risk_level"), + "is_client_work": audit_data.get("is_client_work") + }) + logger.info(f"[{audited + failed}/{len(agents_with_instructions)}] " + f"Classified '{agent.get('agent_name')}' -> Cat {audit_data.get('category')}") + + except Exception as e: + failed += 1 + results.append({"agent_name": agent.get("agent_name"), "error": str(e)}) + logger.error(f"[{audited + failed}/{len(agents_with_instructions)}] " + f"Failed '{agent.get('agent_name')}': {e}") + + # Pause between batches to avoid rate limits (4 seconds per batch) + if i + batch_size < len(agents_with_instructions): + await asyncio.sleep(4) + + logger.info(f"Audit batch complete: {audited} audited, {failed} failed, {skipped} skipped") + + return { + "status": "completed", + "total": total, + "audited_count": audited, + "failed_count": failed, + "skipped_count": skipped, + "results_summary": results[:50] + } + + +async def get_all_audit_results() -> list: + """Get all agents with audit data for the audit tab.""" + agents = await agents_collection.find( + {}, + { + "agent_name": 1, "audit_status": 1, "audit_date": 1, + "audit_category": 1, "audit_risk_level": 1, "audit_summary": 1, + "audit_flags": 1, "audit_recommendations": 1, + "audit_category_reasoning": 1, + "audit_discipline": 1, "audit_discipline_reasoning": 1, + "audit_department": 1, + "audit_is_client_work": 1, "audit_client_work_reasoning": 1, + "audit_client_name_detected": 1, + "discipline": 1, "agent_department": 1, + "client": 1, "client_name": 1, "verification_status": 1, + "instructions": 1, "agent_capabilities": 1, + "created_by": 1, "agent_description": 1, + "audit_reviewer": 1, "audit_reviewer_notes": 1, + "audit_reviewed_date": 1 + } + ).to_list(length=None) + + # Convert ObjectId to string + for agent in agents: + agent["_id"] = str(agent["_id"]) + + return agents + + +async def update_audit_review(agent_id: str, status: str, notes: str, reviewer_email: str): + """Mark an agent's audit as reviewed/cleared.""" + now = datetime.utcnow() + + await agents_collection.update_one( + {"_id": ObjectId(agent_id)}, + {"$set": { + "audit_status": status, + "audit_reviewer": reviewer_email, + "audit_reviewer_notes": notes, + "audit_reviewed_date": now.isoformat() + }} + ) + + # Also record in history + await audit_history_collection.insert_one({ + "agent_id": agent_id, + "audit_date": now, + "action": "review", + "status": status, + "reviewer": reviewer_email, + "notes": notes + }) diff --git a/database.py b/database.py index 2ace175..5ce7e41 100644 --- a/database.py +++ b/database.py @@ -14,6 +14,7 @@ agents_collection = db.get_collection("agents") agent_usage_collection = db.get_collection("agent_usage") notifications_collection = db.get_collection("token_notifications") agent_ratings_collection = db.get_collection("agent_ratings") +audit_history_collection = db.get_collection("audit_history") async def ensure_indexes(): """Create database indexes for performance""" @@ -23,6 +24,8 @@ async def ensure_indexes(): unique=True ) await agents_collection.create_index([("verification_status", 1)]) + await agents_collection.create_index([("audit_status", 1)]) + await audit_history_collection.create_index([("agent_id", 1), ("audit_date", -1)]) print("Database indexes ensured successfully") except Exception as e: print(f"Warning: Failed to create indexes: {e}") diff --git a/main.py b/main.py index d898d41..f098af7 100644 --- a/main.py +++ b/main.py @@ -11,6 +11,7 @@ import auth import config import msal_auth import notifications +import audit_analyzer from datetime import datetime import os import re @@ -243,6 +244,11 @@ def create_agent_response(agent: dict) -> models.AiAgentResponse: verification_status=agent.get("verification_status"), verified_by=agent.get("verified_by"), verified_date=agent.get("verified_date"), + instructions=agent.get("instructions"), + audit_status=agent.get("audit_status"), + audit_date=agent.get("audit_date"), + audit_category=agent.get("audit_category"), + audit_risk_level=agent.get("audit_risk_level"), created_by=agent["created_by"], # Usage tracking fields usage_timeline=agent.get("usage_timeline"), @@ -308,6 +314,7 @@ def map_agent_collector_to_internal(collector_data: models.AgentCollectorCreate) "client": collector_data.client, "client_name": collector_data.client_name, "studio_name": collector_data.studio_name, + "instructions": collector_data.instructions, # Usage tracking fields "usage_timeline": usage_timeline, "conversation_count": collector_data.conversation_count, @@ -1261,6 +1268,61 @@ async def trigger_weekly_digest(current_user: dict = Depends(require_admin)): except Exception as e: raise HTTPException(status_code=500, detail=f"Failed to send digest: {str(e)}") +# Prompt Audit Endpoints +@app.post("/api/admin/audit/run") +async def run_audit( + request: Request, + current_user: dict = Depends(require_admin) +): + """Run Gemini-based audit on agents. Optional body: {agent_id, unclassified_only}""" + if not audit_analyzer.is_gemini_configured(): + raise HTTPException(status_code=503, detail="GOOGLE_API_KEY not configured") + + body = {} + try: + body = await request.json() + except Exception: + pass + + agent_id = body.get("agent_id") + unclassified_only = body.get("unclassified_only", False) + + result = await audit_analyzer.run_audit_batch( + unclassified_only=unclassified_only, + single_agent_id=agent_id + ) + + if result.get("error"): + raise HTTPException(status_code=503, detail=result["error"]) + + return result + +@app.get("/api/admin/audit/results") +async def get_audit_results(current_user: dict = Depends(require_admin_or_readonly)): + """Get all agents with audit data for the Prompt Audit tab""" + agents = await audit_analyzer.get_all_audit_results() + return { + "agents": agents, + "config_status": { + "gemini_configured": audit_analyzer.is_gemini_configured() + } + } + +@app.put("/api/admin/audit/{agent_id}/review") +async def review_audit( + agent_id: str, + review: models.AuditReviewRequest, + current_user: dict = Depends(require_admin) +): + """Mark an agent's audit as reviewed/cleared""" + await audit_analyzer.update_audit_review( + agent_id=agent_id, + status=review.audit_status, + notes=review.reviewer_notes or "", + reviewer_email=current_user["email"] + ) + return {"message": f"Audit status updated to {review.audit_status}"} + @app.get("/api/admin/analytics") async def get_admin_analytics( status: Optional[str] = None, @@ -1369,7 +1431,8 @@ async def export_agents_csv(current_user: dict = Depends(require_admin)): "completion_tokens", "discipline", "rating", - "rating_count" + "rating_count", + "instructions" ] writer = csv.DictWriter(output, fieldnames=fieldnames) @@ -1409,7 +1472,8 @@ async def export_agents_csv(current_user: dict = Depends(require_admin)): "completion_tokens": str(agent.get("completion_tokens", "")) if agent.get("completion_tokens") is not None else "", "discipline": agent.get("discipline", ""), "rating": str(agent.get("rating", "")) if agent.get("rating") is not None else "", - "rating_count": str(agent.get("rating_count", "")) if agent.get("rating_count") is not None else "" + "rating_count": str(agent.get("rating_count", "")) if agent.get("rating_count") is not None else "", + "instructions": agent.get("instructions", "") } writer.writerow(row) @@ -1615,6 +1679,7 @@ async def create_agent_collector( update_fields = { "url": internal_data.get("url"), "discipline": internal_data.get("discipline"), + "instructions": internal_data.get("instructions"), "usage_timeline": internal_data.get("usage_timeline"), "conversation_count": internal_data.get("conversation_count"), "unique_users": internal_data.get("unique_users"), @@ -1646,6 +1711,10 @@ async def create_agent_collector( except Exception as notify_err: print(f"Notification check failed for '{agent.name}': {notify_err}") + # Auto-classify with Gemini in background (non-blocking) + if audit_analyzer.is_gemini_configured() and internal_data.get("instructions"): + asyncio.create_task(audit_analyzer.classify_single_agent(str(existing_agent["_id"]))) + return models.AgentUsageTrackingResponse( status="usage_logged", message="Agent already exists, usage tracked", @@ -1654,13 +1723,13 @@ async def create_agent_collector( else: # Agent doesn't exist - create new registration internal_data = map_agent_collector_to_internal(agent) - + # Handle datetime fields if provided if agent.creation_date: internal_data["agent_created_at"] = agent.creation_date if agent.last_updated: internal_data["agent_updated_at"] = agent.last_updated - + # Create agent using collector-specific function created_agent = await crud.create_agent_from_collector(internal_data) @@ -1670,6 +1739,10 @@ async def create_agent_collector( except Exception as notify_err: print(f"Notification check failed for '{agent.name}': {notify_err}") + # Auto-classify with Gemini in background (non-blocking) + if audit_analyzer.is_gemini_configured() and internal_data.get("instructions"): + asyncio.create_task(audit_analyzer.classify_single_agent(str(created_agent["_id"]))) + return models.AgentCollectorResponse( status="success", message="Agent data collected successfully", diff --git a/models.py b/models.py index 4eeb951..a0430eb 100644 --- a/models.py +++ b/models.py @@ -34,6 +34,7 @@ class AiAgent(BaseModel): last_edited_by: str | None = Field(default=None, title="Email of user who last edited this agent") discipline: str | None = Field(default=None, title="Business discipline/category", max_length=100) rating: float | None = Field(default=None, title="Star rating (1-5)", ge=1, le=5) + instructions: str | None = Field(default=None, title="System prompt / instructions from LibreChat") @@ -107,6 +108,7 @@ class AiAgentCreate(BaseModel): client: Optional[str] = None client_name: Optional[str] = None studio_name: Optional[str] = None + instructions: Optional[str] = None class AiAgentResponse(BaseModel): agent_id: str @@ -141,6 +143,11 @@ class AiAgentResponse(BaseModel): verification_status: Optional[str] = None verified_by: Optional[str] = None verified_date: Optional[str] = None + instructions: Optional[str] = None + audit_status: Optional[str] = None + audit_date: Optional[str] = None + audit_category: Optional[str] = None + audit_risk_level: Optional[str] = None created_by: str # Usage tracking fields (new) @@ -176,6 +183,7 @@ class AgentCollectorCreate(BaseModel): client: Optional[str] = None client_name: Optional[str] = None studio_name: Optional[str] = None + instructions: Optional[str] = None # Usage tracking fields (new) usage_timeline: Optional[List[UsageTimelineEntry]] = None @@ -210,6 +218,10 @@ class AgentUsageRecord(BaseModel): timestamp: str usage_count: Optional[int] = None +class AuditReviewRequest(BaseModel): + audit_status: str = Field(..., pattern="^(flagged|reviewed|cleared)$") + reviewer_notes: Optional[str] = None + class AgentUsageStatsResponse(BaseModel): agent_name: str total_usage_count: int diff --git a/requirements.txt b/requirements.txt index 352fbe5..2d724d5 100644 --- a/requirements.txt +++ b/requirements.txt @@ -39,3 +39,4 @@ msal==1.26.0 itsdangerous==2.2.0 cryptography==41.0.7 apscheduler>=3.10.0 +google-genai>=1.0.0 diff --git a/templates/admin/dashboard.html b/templates/admin/dashboard.html index 71ee903..d140ddc 100644 --- a/templates/admin/dashboard.html +++ b/templates/admin/dashboard.html @@ -109,6 +109,12 @@ +
  • + + Prompt Audit + + +
    • @@ -387,6 +393,156 @@
    + +
    + + + + +
    +
    Prompt Audit & Classification
    +
    + + +
    +
    + + + + + +
    +
    +
    +
    +
    0
    + Audited +
    +
    +
    +
    +
    +
    +
    0
    + Flagged +
    +
    +
    +
    +
    +
    +
    0
    + Reviewed +
    +
    +
    +
    +
    +
    +
    0
    + Cleared +
    +
    +
    +
    +
    +
    +
    0
    + Client Detected +
    +
    +
    +
    +
    +
    +
    0
    + No Instructions +
    +
    +
    +
    + + +
    + Filters: + + + + + + +
    + + +
    + + + + + + + + + + + + + + + + + + + + +
    Agent NameDisciplineDepartmentCategoryRiskClientFlagsStatusLast AuditedActions
    + Click "Run Full Audit" to classify agents, or loading... +
    +
    +
    + @@ -394,6 +550,34 @@ + + +