mirror of
https://github.com/SpecterOps/Nemesis
synced 2026-06-08 12:36:42 +00:00
Merge pull request #91 from SpecterOps/llm_chat_over_data
Llm chat over data
This commit is contained in:
Will Schroeder
@@ -1049,6 +1049,9 @@ services:
|
||||
- NEMESIS_MONITORING=${NEMESIS_MONITORING:-disabled}
|
||||
- OTEL_EXPORTER_OTLP_TRACES_ENDPOINT=http://otel-collector:4317
|
||||
- OTEL_EXPORTER_OTLP_TRACES_ENDPOINT_INSECURE=true
|
||||
# Chatbot configuration
|
||||
- CHATBOT_DB_PASSWORD=${CHATBOT_DB_PASSWORD:-chatbot_pass_change_me}
|
||||
- MCP_MAX_RESULTS=${MCP_MAX_RESULTS:-1000}
|
||||
logging: *logging-config
|
||||
depends_on:
|
||||
postgres: { condition: service_healthy }
|
||||
|
||||
+13
@@ -65,3 +65,16 @@ NEMESIS_URL="https://localhost:7443/"
|
||||
# Phoenix UI will be available at http://localhost:6006
|
||||
# Example:
|
||||
# PHOENIX_ENABLED=true
|
||||
|
||||
|
||||
# (Optional) Chatbot database configuration.
|
||||
# Password for the read-only database user used by the chatbot.
|
||||
# Example:
|
||||
# CHATBOT_DB_PASSWORD="chatbot_secure_password"
|
||||
CHATBOT_DB_PASSWORD="chatbot_pass_change_me"
|
||||
|
||||
# (Optional) Maximum number of results returned by chatbot tools
|
||||
# Helps prevent expensive queries and context window issues.
|
||||
# Example:
|
||||
# MCP_MAX_RESULTS=500
|
||||
MCP_MAX_RESULTS=1000
|
||||
|
||||
@@ -837,4 +837,36 @@ CREATE OR REPLACE TRIGGER update_dpapi_domain_backup_keys_updated_at
|
||||
CREATE OR REPLACE TRIGGER update_dpapi_system_credentials_updated_at
|
||||
BEFORE UPDATE ON dpapi.system_credentials
|
||||
FOR EACH ROW
|
||||
EXECUTE FUNCTION update_updated_at_column();
|
||||
EXECUTE FUNCTION update_updated_at_column();
|
||||
|
||||
|
||||
-----------------------
|
||||
-- CHATBOT READ-ONLY USER
|
||||
-----------------------
|
||||
-- Create read-only user for chatbot queries with restricted table access
|
||||
DO $$
|
||||
BEGIN
|
||||
IF NOT EXISTS (SELECT FROM pg_catalog.pg_roles WHERE rolname = 'chatbot_readonly') THEN
|
||||
CREATE USER chatbot_readonly WITH PASSWORD 'chatbot_pass_change_me';
|
||||
END IF;
|
||||
END
|
||||
$$;
|
||||
|
||||
-- Grant connection and schema usage
|
||||
GRANT CONNECT ON DATABASE enrichment TO chatbot_readonly;
|
||||
GRANT USAGE ON SCHEMA public TO chatbot_readonly;
|
||||
GRANT USAGE ON SCHEMA chromium TO chatbot_readonly;
|
||||
|
||||
-- Grant SELECT on specific tables only (chatbot-accessible tables)
|
||||
GRANT SELECT ON files_enriched TO chatbot_readonly;
|
||||
GRANT SELECT ON enrichments TO chatbot_readonly;
|
||||
GRANT SELECT ON findings TO chatbot_readonly;
|
||||
GRANT SELECT ON file_linkings TO chatbot_readonly;
|
||||
GRANT SELECT ON chromium.cookies TO chatbot_readonly;
|
||||
GRANT SELECT ON chromium.logins TO chatbot_readonly;
|
||||
|
||||
-- Explicitly revoke write permissions to ensure read-only access
|
||||
REVOKE INSERT, UPDATE, DELETE, TRUNCATE ON ALL TABLES IN SCHEMA public FROM chatbot_readonly;
|
||||
REVOKE INSERT, UPDATE, DELETE, TRUNCATE ON ALL TABLES IN SCHEMA chromium FROM chatbot_readonly;
|
||||
REVOKE CREATE ON SCHEMA public FROM chatbot_readonly;
|
||||
REVOKE CREATE ON SCHEMA chromium FROM chatbot_readonly;
|
||||
@@ -7,7 +7,17 @@ RUN apt-get update && \
|
||||
apt-get install -y libpq5 \
|
||||
gcc libc6-dev curl wget libicu-dev && \
|
||||
apt-get clean && \
|
||||
rm -rf /var/lib/apt/lists/*
|
||||
rm -rf /var/lib/apt/lists/* && \
|
||||
ARCH=$(dpkg --print-architecture) && \
|
||||
wget https://go.dev/dl/go1.25.4.linux-${ARCH}.tar.gz && \
|
||||
tar -C /usr/local -xzf go1.25.4.linux-${ARCH}.tar.gz && \
|
||||
rm go1.25.4.linux-${ARCH}.tar.gz
|
||||
|
||||
|
||||
# Install genai-toolbox for chatbot MCP functionality
|
||||
ENV GOPATH=/go \
|
||||
PATH=/usr/local/go/bin:/go/bin:$PATH
|
||||
RUN go install github.com/googleapis/genai-toolbox@v0.18.0
|
||||
|
||||
|
||||
# Install .NET Runtime
|
||||
@@ -82,6 +92,11 @@ ENV DOTNET_ROOT=/usr/local/dotnet \
|
||||
PATH=/usr/local/dotnet:$PATH
|
||||
COPY --from=base /usr/local/dotnet /usr/local/dotnet
|
||||
|
||||
# Copy genai-toolbox binary from base stage
|
||||
ENV GOPATH=/go \
|
||||
PATH=/go/bin:$PATH
|
||||
COPY --from=base /go/bin/genai-toolbox /go/bin/genai-toolbox
|
||||
|
||||
COPY --from=bundle /venv /venv
|
||||
|
||||
|
||||
|
||||
@@ -677,6 +677,23 @@ def run_report_generator(request: dict):
|
||||
return {"success": False, "error": str(e)}
|
||||
|
||||
|
||||
@app.post("/agents/chatbot/stream")
|
||||
async def chatbot_stream_endpoint(request: dict):
|
||||
"""Stream chatbot responses for interactive querying."""
|
||||
try:
|
||||
from agents.tasks.chatbot import ChatbotRequest, chatbot_stream
|
||||
|
||||
# Parse and validate request
|
||||
chatbot_request = ChatbotRequest(**request)
|
||||
|
||||
# Stream the response
|
||||
return await chatbot_stream(chatbot_request)
|
||||
|
||||
except Exception as e:
|
||||
logger.exception(message="Error in chatbot streaming")
|
||||
return {"success": False, "error": str(e)}
|
||||
|
||||
|
||||
@app.api_route("/healthz", methods=["GET", "HEAD"])
|
||||
async def health_check():
|
||||
"""Health check endpoint."""
|
||||
|
||||
@@ -0,0 +1,506 @@
|
||||
# genai-toolbox configuration for Nemesis Chatbot
|
||||
# Database source definition
|
||||
sources:
|
||||
chatbot-db:
|
||||
kind: postgres
|
||||
host: ${POSTGRES_HOST:postgres}
|
||||
port: ${POSTGRES_PORT:5432}
|
||||
database: ${POSTGRES_DB:enrichment}
|
||||
user: chatbot_readonly
|
||||
password: ${CHATBOT_DB_PASSWORD}
|
||||
|
||||
# Custom SQL tools for querying Nemesis data
|
||||
tools:
|
||||
# FILES_ENRICHED queries
|
||||
count-files:
|
||||
kind: postgres-sql
|
||||
source: chatbot-db
|
||||
description: Count total files, optionally filtered by project, agent, source, or extension
|
||||
parameters:
|
||||
- name: project
|
||||
type: string
|
||||
required: false
|
||||
description: Filter by project name (optional)
|
||||
- name: agent_id
|
||||
type: string
|
||||
required: false
|
||||
description: Filter by agent ID (optional)
|
||||
- name: source
|
||||
type: string
|
||||
required: false
|
||||
description: Filter by source (case-insensitive pattern match, optional)
|
||||
- name: extension
|
||||
type: string
|
||||
required: false
|
||||
description: Filter by file extension (optional)
|
||||
statement: |
|
||||
SELECT COUNT(*) as file_count
|
||||
FROM files_enriched
|
||||
WHERE ($1::text IS NULL OR project = $1)
|
||||
AND ($2::text IS NULL OR agent_id = $2)
|
||||
AND ($3::text IS NULL OR LOWER(source) LIKE LOWER('%' || $3 || '%'))
|
||||
AND ($4::text IS NULL OR extension = $4)
|
||||
|
||||
search-files:
|
||||
kind: postgres-sql
|
||||
source: chatbot-db
|
||||
description: Search for files by name, path, or extension with optional filters
|
||||
parameters:
|
||||
- name: filename_pattern
|
||||
type: string
|
||||
required: false
|
||||
description: Search pattern for filename (case-insensitive, optional)
|
||||
- name: path_pattern
|
||||
type: string
|
||||
required: false
|
||||
description: Search pattern for file path (case-insensitive, optional)
|
||||
- name: extension
|
||||
type: string
|
||||
required: false
|
||||
description: Filter by file extension (optional)
|
||||
- name: project
|
||||
type: string
|
||||
required: false
|
||||
description: Filter by project name (optional)
|
||||
- name: source
|
||||
type: string
|
||||
required: false
|
||||
description: Filter by source (case-insensitive pattern match, optional)
|
||||
- name: limit
|
||||
type: integer
|
||||
required: false
|
||||
description: Maximum number of results (default 100, max 1000)
|
||||
statement: |
|
||||
SELECT object_id::text, file_name, path, extension, size, magic_type, mime_type,
|
||||
source, agent_id, project, timestamp, originating_object_id::text
|
||||
FROM files_enriched
|
||||
WHERE ($1::text IS NULL OR LOWER(file_name) LIKE LOWER('%' || $1 || '%'))
|
||||
AND ($2::text IS NULL OR LOWER(path) LIKE LOWER('%' || $2 || '%'))
|
||||
AND ($3::text IS NULL OR extension = $3)
|
||||
AND ($4::text IS NULL OR project = $4)
|
||||
AND ($5::text IS NULL OR LOWER(source) LIKE LOWER('%' || $5 || '%'))
|
||||
ORDER BY timestamp DESC
|
||||
LIMIT LEAST(COALESCE($6, 100), 1000)
|
||||
|
||||
get-file-details:
|
||||
kind: postgres-sql
|
||||
source: chatbot-db
|
||||
description: Get detailed information about a specific file by object_id
|
||||
parameters:
|
||||
- name: object_id
|
||||
type: string
|
||||
description: The UUID of the file to retrieve
|
||||
statement: |
|
||||
SELECT object_id::text, agent_id, source, project, timestamp, path, file_name,
|
||||
extension, size, magic_type, mime_type, originating_object_id::text
|
||||
FROM files_enriched
|
||||
WHERE object_id = $1::uuid
|
||||
|
||||
# ENRICHMENTS queries
|
||||
list-enrichment-modules:
|
||||
kind: postgres-sql
|
||||
source: chatbot-db
|
||||
description: List all unique enrichment module names with their usage counts
|
||||
statement: |
|
||||
SELECT module_name, COUNT(*) as usage_count
|
||||
FROM enrichments
|
||||
GROUP BY module_name
|
||||
ORDER BY usage_count DESC
|
||||
|
||||
get-file-enrichments:
|
||||
kind: postgres-sql
|
||||
source: chatbot-db
|
||||
description: Get all enrichment results for a specific file
|
||||
parameters:
|
||||
- name: object_id
|
||||
type: string
|
||||
description: The UUID of the file
|
||||
statement: |
|
||||
SELECT e.enrichment_id, e.object_id::text, e.module_name, e.result_data, e.created_at,
|
||||
f.file_name, f.path
|
||||
FROM enrichments e
|
||||
JOIN files_enriched f ON e.object_id = f.object_id
|
||||
WHERE e.object_id = $1::uuid
|
||||
ORDER BY e.created_at DESC
|
||||
|
||||
search-enrichments-by-module:
|
||||
kind: postgres-sql
|
||||
source: chatbot-db
|
||||
description: Search enrichment results from a specific module
|
||||
parameters:
|
||||
- name: module_name
|
||||
type: string
|
||||
description: Name of the enrichment module
|
||||
- name: limit
|
||||
type: integer
|
||||
required: false
|
||||
description: Maximum number of results (default 100, max 1000)
|
||||
statement: |
|
||||
SELECT e.object_id::text, e.module_name, e.result_data, e.created_at,
|
||||
f.file_name, f.path, f.source
|
||||
FROM enrichments e
|
||||
JOIN files_enriched f ON e.object_id = f.object_id
|
||||
WHERE e.module_name = $1
|
||||
ORDER BY e.created_at DESC
|
||||
LIMIT LEAST(COALESCE($2, 100), 1000)
|
||||
|
||||
# FINDINGS queries
|
||||
get-unique-findings-categories:
|
||||
kind: postgres-sql
|
||||
source: chatbot-db
|
||||
description: List all unique finding categories with their counts and severity statistics
|
||||
statement: |
|
||||
SELECT category, COUNT(*) as count,
|
||||
AVG(severity) as avg_severity,
|
||||
MAX(severity) as max_severity
|
||||
FROM findings
|
||||
GROUP BY category
|
||||
ORDER BY count DESC
|
||||
|
||||
count-findings:
|
||||
kind: postgres-sql
|
||||
source: chatbot-db
|
||||
description: Count findings, optionally filtered by severity, category, or source
|
||||
parameters:
|
||||
- name: min_severity
|
||||
type: integer
|
||||
required: false
|
||||
description: Minimum severity level (0-10, optional)
|
||||
- name: category
|
||||
type: string
|
||||
required: false
|
||||
description: Filter by finding category (optional)
|
||||
- name: source
|
||||
type: string
|
||||
required: false
|
||||
description: Filter by source (case-insensitive pattern match, optional)
|
||||
statement: |
|
||||
SELECT COUNT(*) as finding_count
|
||||
FROM findings f
|
||||
JOIN files_enriched fe ON f.object_id = fe.object_id
|
||||
WHERE ($1::integer IS NULL OR f.severity >= $1)
|
||||
AND ($2::text IS NULL OR $2 = '' OR f.category = $2)
|
||||
AND ($3::text IS NULL OR $3 = '' OR LOWER(fe.source) LIKE LOWER('%' || $3 || '%'))
|
||||
|
||||
search-findings:
|
||||
kind: postgres-sql
|
||||
source: chatbot-db
|
||||
description: Search findings with filters for severity, category, and source
|
||||
parameters:
|
||||
- name: min_severity
|
||||
type: integer
|
||||
required: false
|
||||
description: Minimum severity level (0-10, optional)
|
||||
- name: category
|
||||
type: string
|
||||
required: false
|
||||
description: Filter by finding category (optional)
|
||||
- name: source
|
||||
type: string
|
||||
required: false
|
||||
description: Filter by source (case-insensitive pattern match, optional)
|
||||
- name: limit
|
||||
type: integer
|
||||
required: false
|
||||
description: Maximum number of results (default 100, max 1000)
|
||||
statement: |
|
||||
SELECT f.finding_id, f.finding_name, f.category, f.severity, f.origin_type,
|
||||
f.origin_name, f.data, f.created_at,
|
||||
fe.file_name, fe.path, fe.source, fe.agent_id, fe.project
|
||||
FROM findings f
|
||||
JOIN files_enriched fe ON f.object_id = fe.object_id
|
||||
WHERE ($1::integer IS NULL OR f.severity >= $1)
|
||||
AND ($2::text IS NULL OR f.category = $2)
|
||||
AND ($3::text IS NULL OR LOWER(fe.source) LIKE LOWER('%' || $3 || '%'))
|
||||
ORDER BY f.severity DESC, f.created_at DESC
|
||||
LIMIT LEAST(COALESCE($4, 100), 1000)
|
||||
|
||||
get-file-findings:
|
||||
kind: postgres-sql
|
||||
source: chatbot-db
|
||||
description: Get all findings for a specific file by object_id
|
||||
parameters:
|
||||
- name: object_id
|
||||
type: string
|
||||
description: The UUID of the file
|
||||
statement: |
|
||||
SELECT f.finding_id, f.finding_name, f.category, f.severity,
|
||||
f.origin_type, f.origin_name, f.data, f.created_at,
|
||||
fe.file_name, fe.path, fe.source, fe.agent_id, fe.project
|
||||
FROM findings f
|
||||
JOIN files_enriched fe ON f.object_id = fe.object_id
|
||||
WHERE f.object_id = $1::uuid
|
||||
ORDER BY f.severity DESC, f.created_at DESC
|
||||
|
||||
get-findings-by-category:
|
||||
kind: postgres-sql
|
||||
source: chatbot-db
|
||||
description: Get aggregated count of findings grouped by category
|
||||
parameters:
|
||||
- name: source
|
||||
type: string
|
||||
required: false
|
||||
description: Filter by source (case-insensitive pattern match, optional)
|
||||
statement: |
|
||||
SELECT f.category, COUNT(*) as count,
|
||||
AVG(f.severity) as avg_severity,
|
||||
MAX(f.severity) as max_severity
|
||||
FROM findings f
|
||||
JOIN files_enriched fe ON f.object_id = fe.object_id
|
||||
WHERE ($1::text IS NULL OR LOWER(fe.source) LIKE LOWER('%' || $1 || '%'))
|
||||
GROUP BY f.category
|
||||
ORDER BY count DESC
|
||||
|
||||
get-triaged-findings:
|
||||
kind: postgres-sql
|
||||
source: chatbot-db
|
||||
description: Get findings with triage status (true positive, false positive, etc)
|
||||
parameters:
|
||||
- name: triage_value
|
||||
type: string
|
||||
required: false
|
||||
description: Filter by triage value (e.g., true_positive, false_positive, optional)
|
||||
- name: source
|
||||
type: string
|
||||
required: false
|
||||
description: Filter by source (case-insensitive pattern match, optional)
|
||||
- name: limit
|
||||
type: integer
|
||||
required: false
|
||||
description: Maximum number of results (default 100, max 1000)
|
||||
statement: |
|
||||
SELECT f.finding_id, f.finding_name, f.category, f.severity,
|
||||
fth.value as triage_value, fth.explanation, fth.timestamp as triage_timestamp,
|
||||
fe.file_name, fe.path, fe.source
|
||||
FROM findings f
|
||||
JOIN findings_triage_history fth ON f.triage_id = fth.id
|
||||
JOIN files_enriched fe ON f.object_id = fe.object_id
|
||||
WHERE ($1::text IS NULL OR fth.value = $1)
|
||||
AND ($2::text IS NULL OR LOWER(fe.source) LIKE LOWER('%' || $2 || '%'))
|
||||
ORDER BY fth.timestamp DESC
|
||||
LIMIT LEAST(COALESCE($3, 100), 1000)
|
||||
|
||||
search-credential-findings-by-host:
|
||||
kind: postgres-sql
|
||||
source: chatbot-db
|
||||
description: Search for credential findings related to a specific hostname or system. Use this to find credentials that may provide access to a target system.
|
||||
parameters:
|
||||
- name: hostname
|
||||
type: string
|
||||
description: Target hostname or system name to search for in credential findings (case-insensitive partial match)
|
||||
- name: source
|
||||
type: string
|
||||
required: false
|
||||
description: Filter by source (case-insensitive pattern match, optional)
|
||||
- name: limit
|
||||
type: integer
|
||||
required: false
|
||||
description: Maximum number of results (default 100, max 1000)
|
||||
statement: |
|
||||
SELECT f.finding_id, f.finding_name, f.category, f.severity,
|
||||
f.origin_type, f.origin_name, f.data, f.created_at,
|
||||
fe.file_name, fe.path, fe.source, fe.agent_id, fe.project
|
||||
FROM findings f
|
||||
JOIN files_enriched fe ON f.object_id = fe.object_id
|
||||
WHERE f.category = 'credential'
|
||||
AND LOWER(f.data::text) LIKE LOWER('%' || $1 || '%')
|
||||
AND ($2::text IS NULL OR LOWER(fe.source) LIKE LOWER('%' || $2 || '%'))
|
||||
ORDER BY f.severity DESC, f.created_at DESC
|
||||
LIMIT LEAST(COALESCE($3, 100), 1000)
|
||||
|
||||
search-logins-by-host:
|
||||
kind: postgres-sql
|
||||
source: chatbot-db
|
||||
description: Search for decrypted browser credentials related to a specific hostname or URL. Use this to find saved browser passwords for accessing a target system.
|
||||
parameters:
|
||||
- name: hostname
|
||||
type: string
|
||||
description: Target hostname or URL to search for in saved browser credentials (case-insensitive partial match)
|
||||
- name: source
|
||||
type: string
|
||||
required: false
|
||||
description: Filter by source (case-insensitive pattern match, optional)
|
||||
- name: limit
|
||||
type: integer
|
||||
required: false
|
||||
description: Maximum number of results (default 100, max 1000)
|
||||
statement: |
|
||||
SELECT origin_url, username_value, password_value_dec, signon_realm,
|
||||
date_created, date_last_used, times_used,
|
||||
source, username, browser, agent_id, project
|
||||
FROM chromium.logins
|
||||
WHERE is_decrypted = true
|
||||
AND (LOWER(origin_url) LIKE LOWER('%' || $1 || '%')
|
||||
OR LOWER(signon_realm) LIKE LOWER('%' || $1 || '%'))
|
||||
AND ($2::text IS NULL OR LOWER(source) LIKE LOWER('%' || $2 || '%'))
|
||||
ORDER BY date_last_used DESC NULLS LAST, times_used DESC
|
||||
LIMIT LEAST(COALESCE($3, 100), 1000)
|
||||
|
||||
# CHROMIUM.COOKIES queries
|
||||
count-cookies:
|
||||
kind: postgres-sql
|
||||
source: chatbot-db
|
||||
description: Count browser cookies, optionally filtered by host or source
|
||||
parameters:
|
||||
- name: host_pattern
|
||||
type: string
|
||||
required: false
|
||||
description: Filter by host_key pattern (case-insensitive, optional)
|
||||
- name: source
|
||||
type: string
|
||||
required: false
|
||||
description: Filter by source (case-insensitive pattern match, optional)
|
||||
statement: |
|
||||
SELECT COUNT(*) as cookie_count
|
||||
FROM chromium.cookies
|
||||
WHERE ($1::text IS NULL OR LOWER(host_key) LIKE LOWER('%' || $1 || '%'))
|
||||
AND ($2::text IS NULL OR LOWER(source) LIKE LOWER('%' || $2 || '%'))
|
||||
|
||||
search-cookies:
|
||||
kind: postgres-sql
|
||||
source: chatbot-db
|
||||
description: Search browser cookies by host, name, or source
|
||||
parameters:
|
||||
- name: host_pattern
|
||||
type: string
|
||||
required: false
|
||||
description: Filter by host_key pattern (case-insensitive, optional)
|
||||
- name: name_pattern
|
||||
type: string
|
||||
required: false
|
||||
description: Filter by cookie name pattern (case-insensitive, optional)
|
||||
- name: source
|
||||
type: string
|
||||
required: false
|
||||
description: Filter by source (case-insensitive pattern match, optional)
|
||||
- name: is_decrypted
|
||||
type: boolean
|
||||
required: false
|
||||
description: Filter by decryption status (optional)
|
||||
- name: limit
|
||||
type: integer
|
||||
required: false
|
||||
description: Maximum number of results (default 100, max 1000)
|
||||
statement: |
|
||||
SELECT host_key, name, path, value_dec, is_decrypted, is_secure, is_httponly,
|
||||
expires_utc, source, username, browser, agent_id, project, originating_object_id::text
|
||||
FROM chromium.cookies
|
||||
WHERE ($1::text IS NULL OR LOWER(host_key) LIKE LOWER('%' || $1 || '%'))
|
||||
AND ($2::text IS NULL OR LOWER(name) LIKE LOWER('%' || $2 || '%'))
|
||||
AND ($3::text IS NULL OR LOWER(source) LIKE LOWER('%' || $3 || '%'))
|
||||
AND ($4::boolean IS NULL OR is_decrypted = $4)
|
||||
ORDER BY expires_utc DESC NULLS LAST
|
||||
LIMIT LEAST(COALESCE($5, 100), 1000)
|
||||
|
||||
# CHROMIUM.LOGINS queries
|
||||
count-logins:
|
||||
kind: postgres-sql
|
||||
source: chatbot-db
|
||||
description: Count saved browser credentials, optionally filtered by URL or source
|
||||
parameters:
|
||||
- name: url_pattern
|
||||
type: string
|
||||
required: false
|
||||
description: Filter by origin_url pattern (case-insensitive, optional)
|
||||
- name: source
|
||||
type: string
|
||||
required: false
|
||||
description: Filter by source (case-insensitive pattern match, optional)
|
||||
- name: is_decrypted
|
||||
type: boolean
|
||||
required: false
|
||||
description: Filter by decryption status (optional)
|
||||
statement: |
|
||||
SELECT COUNT(*) as login_count
|
||||
FROM chromium.logins
|
||||
WHERE ($1::text IS NULL OR LOWER(origin_url) LIKE LOWER('%' || $1 || '%'))
|
||||
AND ($2::text IS NULL OR LOWER(source) LIKE LOWER('%' || $2 || '%'))
|
||||
AND ($3::boolean IS NULL OR is_decrypted = $3)
|
||||
|
||||
search-logins:
|
||||
kind: postgres-sql
|
||||
source: chatbot-db
|
||||
description: Search browser saved credentials by URL, username, or source
|
||||
parameters:
|
||||
- name: url_pattern
|
||||
type: string
|
||||
required: false
|
||||
description: Filter by origin_url pattern (case-insensitive, optional)
|
||||
- name: username_pattern
|
||||
type: string
|
||||
required: false
|
||||
description: Filter by username_value pattern (case-insensitive, optional)
|
||||
- name: source
|
||||
type: string
|
||||
required: false
|
||||
description: Filter by source (case-insensitive pattern match, optional)
|
||||
- name: is_decrypted
|
||||
type: boolean
|
||||
required: false
|
||||
description: Filter by decryption status (optional)
|
||||
- name: limit
|
||||
type: integer
|
||||
required: false
|
||||
description: Maximum number of results (default 100, max 1000)
|
||||
statement: |
|
||||
SELECT origin_url, username_value, password_value_dec, signon_realm,
|
||||
is_decrypted, date_created, date_last_used, times_used,
|
||||
source, username, browser, agent_id, project, originating_object_id::text
|
||||
FROM chromium.logins
|
||||
WHERE ($1::text IS NULL OR LOWER(origin_url) LIKE LOWER('%' || $1 || '%'))
|
||||
AND ($2::text IS NULL OR LOWER(username_value) LIKE LOWER('%' || $2 || '%'))
|
||||
AND ($3::text IS NULL OR LOWER(source) LIKE LOWER('%' || $3 || '%'))
|
||||
AND ($4::boolean IS NULL OR is_decrypted = $4)
|
||||
ORDER BY date_last_used DESC NULLS LAST
|
||||
LIMIT LEAST(COALESCE($5, 100), 1000)
|
||||
|
||||
# PLAINTEXT CONTENT searches
|
||||
search-document-content:
|
||||
kind: postgres-sql
|
||||
source: chatbot-db
|
||||
description: Full-text search through plaintext content chunks of files. Returns the first matching chunk from each file.
|
||||
parameters:
|
||||
- name: search_query
|
||||
type: string
|
||||
description: Text to search for in document content (full-text search)
|
||||
- name: path_pattern
|
||||
type: string
|
||||
required: false
|
||||
description: Filter by file path pattern using LIKE (e.g., '%folder%', optional)
|
||||
- name: agent_pattern
|
||||
type: string
|
||||
required: false
|
||||
description: Filter by agent ID pattern using LIKE (optional)
|
||||
- name: project_name
|
||||
type: string
|
||||
required: false
|
||||
description: Filter by exact project name (optional)
|
||||
- name: start_date
|
||||
type: string
|
||||
required: false
|
||||
description: Filter files from this date onwards in ISO format (e.g., '2024-01-01', optional)
|
||||
- name: end_date
|
||||
type: string
|
||||
required: false
|
||||
description: Filter files up to this date in ISO format (e.g., '2024-12-31', optional)
|
||||
- name: max_results
|
||||
type: integer
|
||||
required: false
|
||||
description: Maximum number of results (default 100, max 1000)
|
||||
- name: source_pattern
|
||||
type: string
|
||||
required: false
|
||||
description: Filter by source pattern using LIKE (optional)
|
||||
statement: |
|
||||
SELECT object_id::text, chunk_number, content, file_name, path, extension,
|
||||
project, agent_id, source, timestamp
|
||||
FROM public.search_documents(
|
||||
$1,
|
||||
$2,
|
||||
$3,
|
||||
$4,
|
||||
$5::timestamp with time zone,
|
||||
$6::timestamp with time zone,
|
||||
COALESCE($7, 100),
|
||||
$8
|
||||
)
|
||||
@@ -0,0 +1,355 @@
|
||||
"""Chatbot agent for interactive querying of Nemesis data."""
|
||||
|
||||
import asyncio
|
||||
import os
|
||||
import subprocess
|
||||
from pathlib import Path
|
||||
from typing import AsyncGenerator
|
||||
|
||||
import structlog
|
||||
from agents.base_agent import BaseAgent
|
||||
from agents.logger import set_agent_metadata
|
||||
from agents.model_manager import ModelManager
|
||||
from agents.prompt_manager import PromptManager
|
||||
from common.db import get_postgres_connection_str
|
||||
from fastapi import HTTPException
|
||||
from fastapi.responses import StreamingResponse
|
||||
from pydantic import BaseModel
|
||||
from pydantic_ai import Agent
|
||||
from pydantic_ai.mcp import MCPServerStreamableHTTP
|
||||
from pydantic_ai.settings import ModelSettings
|
||||
|
||||
logger = structlog.get_logger(__name__)
|
||||
|
||||
|
||||
class ChatMessage(BaseModel):
|
||||
"""A single chat message."""
|
||||
|
||||
role: str # "user" or "assistant"
|
||||
content: str
|
||||
|
||||
|
||||
class ChatbotRequest(BaseModel):
|
||||
"""Request model for chatbot queries."""
|
||||
|
||||
message: str
|
||||
history: list[ChatMessage] = []
|
||||
use_history: bool = True
|
||||
temperature: float = 0.7
|
||||
|
||||
|
||||
class ChatbotAgent(BaseAgent):
|
||||
"""Agent for interactive data querying via natural language."""
|
||||
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
self.prompt_manager = PromptManager(get_postgres_connection_str())
|
||||
self.name = "chatbot"
|
||||
self.description = "Interactive chatbot for querying Nemesis data"
|
||||
self.agent_type = "llm_based"
|
||||
self.has_prompt = True
|
||||
self.llm_temperature = 0.7 # Default, can be overridden per request
|
||||
|
||||
# Get max rows from environment
|
||||
max_results = int(os.getenv("MCP_MAX_RESULTS", "1000"))
|
||||
|
||||
# System prompt - will be saved to DB on first use
|
||||
self.system_prompt = f"""You are a data query assistant for Nemesis, an offensive security data platform.
|
||||
|
||||
Your role is to retrieve and report data from the database. Do NOT provide recommendations, analysis, or suggestions - only report the requested data. You have access to MCP tools to query the Nemesis PostgreSQL database.
|
||||
|
||||
When answering questions:
|
||||
1. Query the database using the appropriate tools
|
||||
2. Report ONLY the data retrieved unless explicitly instructed otherwise - i.e., no analysis or recommendations unless a user explicitly asks for it
|
||||
3. Present results clearly and concisely
|
||||
4. For large result sets, summarize counts and key details
|
||||
5. Use case-insensitive pattern matching for host/source filters
|
||||
6. Be brief - users want facts, not explanations unless they explicitly request them
|
||||
7. Don't return the "project" field to users
|
||||
|
||||
Query Guidelines:
|
||||
- Use a `limit` of {max_results} for maximum results
|
||||
- Filter by severity, category, or source to narrow results
|
||||
- Use search-document-content when users ask to search "for" or "containing" specific text but otherwise restrict your usage of "search-document-content" since it can return a lot of results
|
||||
- a `originating_object_id` field points to the `object_id` the finding/file originated from
|
||||
|
||||
Searching for Credentials to Access Systems:
|
||||
When users ask about accessing a specific system or finding credentials for a hostname, follow this order:
|
||||
1. First, use search-credential-findings-by-host to find credential findings related to the target hostname
|
||||
2. Second, use search-logins-by-host to find decrypted browser credentials for the target hostname
|
||||
3. Only as a last resort, if the above return no results, use search-document-content with the hostname to search file contents
|
||||
This order ensures you check the most relevant credential sources first before falling back to broader document searches.
|
||||
"""
|
||||
|
||||
self.mcp_process = None
|
||||
|
||||
def get_prompt(self) -> str:
|
||||
"""Get the chatbot prompt from database or use default."""
|
||||
try:
|
||||
prompt_data = self.prompt_manager.get_prompt(self.name)
|
||||
|
||||
if prompt_data:
|
||||
return prompt_data["prompt"]
|
||||
else:
|
||||
logger.info("No prompt found in database, initializing with default", agent_name=self.name)
|
||||
success = self.prompt_manager.save_prompt(self.name, self.system_prompt, self.description)
|
||||
if success:
|
||||
logger.info("Default prompt saved to database", agent_name=self.name)
|
||||
else:
|
||||
logger.debug(
|
||||
"Could not save default prompt to database (likely during startup)", agent_name=self.name
|
||||
)
|
||||
|
||||
return self.system_prompt
|
||||
|
||||
except Exception as e:
|
||||
logger.warning("Error managing prompt, using default", agent_name=self.name, error=str(e))
|
||||
return self.system_prompt
|
||||
|
||||
def execute(self, ctx, activity_input: dict) -> dict:
|
||||
"""
|
||||
Execute method required by BaseAgent.
|
||||
|
||||
Note: ChatbotAgent is designed for interactive HTTP streaming,
|
||||
not workflow-based execution. Use the chatbot_stream endpoint instead.
|
||||
"""
|
||||
logger.warning("ChatbotAgent.execute called but this agent is designed for HTTP streaming only")
|
||||
return {
|
||||
"success": False,
|
||||
"error": "ChatbotAgent does not support workflow execution. Use /agents/chatbot/stream endpoint instead."
|
||||
}
|
||||
|
||||
def _get_chatbot_connection_string(self) -> str:
|
||||
"""Get PostgreSQL connection string for chatbot read-only user."""
|
||||
chatbot_password = os.getenv("CHATBOT_DB_PASSWORD", "chatbot_pass_change_me")
|
||||
postgres_host = os.getenv("POSTGRES_HOST", "postgres")
|
||||
postgres_port = os.getenv("POSTGRES_PORT", "5432")
|
||||
postgres_db = os.getenv("POSTGRES_DB", "enrichment")
|
||||
postgres_params = os.getenv("POSTGRES_PARAMETERS", "sslmode=disable")
|
||||
|
||||
return f"postgresql://chatbot_readonly:{chatbot_password}@{postgres_host}:{postgres_port}/{postgres_db}?{postgres_params}"
|
||||
|
||||
async def start_mcp_server(self):
|
||||
"""Start the genai-toolbox MCP server as a subprocess listening on HTTP."""
|
||||
# Check if process is already running
|
||||
if self.mcp_process and self.mcp_process.poll() is None:
|
||||
logger.debug("MCP server already running")
|
||||
return
|
||||
|
||||
try:
|
||||
tools_file = Path(__file__).parent.parent / "mcp" / "tools.yaml"
|
||||
if not tools_file.exists():
|
||||
raise FileNotFoundError(f"tools.yaml not found at {tools_file}")
|
||||
|
||||
# Get database connection string for chatbot readonly user
|
||||
db_url = self._get_chatbot_connection_string()
|
||||
|
||||
# Start genai-toolbox HTTP server (default port 5000)
|
||||
logger.info("Starting genai-toolbox MCP HTTP server", tools_file=str(tools_file))
|
||||
|
||||
self.mcp_process = subprocess.Popen(
|
||||
["genai-toolbox", "--tools-file", str(tools_file)],
|
||||
env={**os.environ, "DATABASE_URL": db_url},
|
||||
stdout=subprocess.PIPE,
|
||||
stderr=subprocess.PIPE,
|
||||
)
|
||||
|
||||
# Give it a moment to start
|
||||
await asyncio.sleep(2)
|
||||
|
||||
if self.mcp_process.poll() is not None:
|
||||
stderr = self.mcp_process.stderr.read().decode() if self.mcp_process.stderr else ""
|
||||
# If it failed due to address in use, that's actually okay
|
||||
if "address already in use" in stderr.lower():
|
||||
logger.info("MCP server already running (address in use)")
|
||||
self.mcp_process = None
|
||||
return
|
||||
raise RuntimeError(f"MCP server failed to start: {stderr}")
|
||||
|
||||
logger.info("MCP HTTP server started successfully on http://127.0.0.1:5000/mcp")
|
||||
|
||||
except Exception as e:
|
||||
logger.error("Failed to start MCP server", error=str(e))
|
||||
raise
|
||||
|
||||
async def stop_mcp_server(self):
|
||||
"""Stop the MCP server subprocess."""
|
||||
if self.mcp_process and self.mcp_process.poll() is None:
|
||||
logger.info("Stopping MCP server")
|
||||
self.mcp_process.terminate()
|
||||
try:
|
||||
self.mcp_process.wait(timeout=5)
|
||||
except subprocess.TimeoutExpired:
|
||||
logger.warning("MCP server didn't stop gracefully, killing")
|
||||
self.mcp_process.kill()
|
||||
self.mcp_process = None
|
||||
|
||||
async def stream_chat_response(self, request: ChatbotRequest) -> AsyncGenerator[str, None]:
|
||||
"""
|
||||
Stream chatbot responses token-by-token.
|
||||
|
||||
Args:
|
||||
request: ChatbotRequest with message, history, and settings
|
||||
|
||||
Yields:
|
||||
Chunks of the response as they're generated
|
||||
"""
|
||||
model = ModelManager.get_model()
|
||||
|
||||
if not model:
|
||||
logger.warning("No model available from ModelManager")
|
||||
raise HTTPException(status_code=503, detail="AI model not available")
|
||||
|
||||
try:
|
||||
# Set metadata for Phoenix tracing
|
||||
set_agent_metadata(
|
||||
agent_name="chatbot",
|
||||
message_length=len(request.message),
|
||||
has_history=len(request.history) > 0 if request.use_history else False,
|
||||
tags=["chatbot", "interactive_query"],
|
||||
)
|
||||
|
||||
# Get current prompt from database
|
||||
current_prompt = self.get_prompt()
|
||||
|
||||
# Build conversation history if enabled
|
||||
conversation = ""
|
||||
if request.use_history and request.history:
|
||||
for msg in request.history:
|
||||
role_label = "User" if msg.role == "user" else "Assistant"
|
||||
conversation += f"{role_label}: {msg.content}\n\n"
|
||||
|
||||
# Add current message
|
||||
conversation += f"User: {request.message}\n\nAssistant:"
|
||||
|
||||
# Connect to MCP HTTP server (genai-toolbox running on http://127.0.0.1:5000/mcp)
|
||||
mcp_server = MCPServerStreamableHTTP(url='http://127.0.0.1:5000/mcp')
|
||||
|
||||
# Create agent with MCP tools
|
||||
agent = Agent(
|
||||
model=model,
|
||||
system_prompt=current_prompt,
|
||||
toolsets=[mcp_server],
|
||||
instrument=ModelManager.is_instrumentation_enabled(),
|
||||
retries=5, # Increased from 2 to handle transient MCP tool failures
|
||||
model_settings=ModelSettings(temperature=request.temperature),
|
||||
)
|
||||
|
||||
logger.info("Starting chatbot stream", message=request.message, temperature=request.temperature)
|
||||
|
||||
# When tools are involved, streaming doesn't work as expected
|
||||
# Get the complete result and send it
|
||||
result = await agent.run(conversation)
|
||||
|
||||
# Log tool calls and their results with full details
|
||||
tool_calls = []
|
||||
tool_errors = []
|
||||
if hasattr(result, 'all_messages'):
|
||||
for msg_idx, msg in enumerate(result.all_messages()):
|
||||
logger.debug(f"Message {msg_idx}: type={type(msg).__name__}, role={getattr(msg, 'role', 'unknown')}")
|
||||
if hasattr(msg, 'parts'):
|
||||
for part_idx, part in enumerate(msg.parts):
|
||||
part_type = type(part).__name__
|
||||
logger.debug(f" Part {part_idx}: type={part_type}")
|
||||
|
||||
# Log tool calls (requests)
|
||||
if hasattr(part, 'tool_name'):
|
||||
tool_info = {
|
||||
'tool': part.tool_name,
|
||||
'args': getattr(part, 'args', {})
|
||||
}
|
||||
logger.info(f"TOOL CALL: {part.tool_name}", args=tool_info['args'])
|
||||
tool_calls.append(tool_info)
|
||||
|
||||
# Log tool returns (responses)
|
||||
if hasattr(part, 'tool_name') and hasattr(part, 'content'):
|
||||
full_content = str(part.content)
|
||||
logger.debug(
|
||||
f"TOOL RESPONSE: {part.tool_name}",
|
||||
content_length=len(full_content),
|
||||
full_content=full_content # Log FULL content for debugging
|
||||
)
|
||||
|
||||
# Log errors
|
||||
if hasattr(part, 'error'):
|
||||
error_info = {
|
||||
'tool': getattr(part, 'tool_name', 'unknown'),
|
||||
'error': str(part.error)
|
||||
}
|
||||
logger.error("TOOL ERROR", error_info=error_info)
|
||||
tool_errors.append(error_info)
|
||||
|
||||
if tool_calls:
|
||||
logger.info("MCP tools called", count=len(tool_calls), tools=[t['tool'] for t in tool_calls])
|
||||
else:
|
||||
logger.warning("No MCP tools were called by the LLM")
|
||||
|
||||
if tool_errors:
|
||||
logger.error("Tool errors occurred", error_count=len(tool_errors))
|
||||
|
||||
# Extract just the text output from the result
|
||||
if hasattr(result, 'data'):
|
||||
final_text = str(result.data)
|
||||
elif hasattr(result, 'output'):
|
||||
final_text = str(result.output)
|
||||
else:
|
||||
final_text = str(result)
|
||||
|
||||
logger.info(f"Got complete response, {len(final_text)} chars")
|
||||
|
||||
# Debug: Log final response to check for UUID corruption
|
||||
if 'object_id' in final_text.lower() or 'uuid' in final_text.lower():
|
||||
logger.warning(
|
||||
"FINAL RESPONSE contains object_id/UUID",
|
||||
final_response=final_text[:2000] # Log first 2000 chars
|
||||
)
|
||||
|
||||
# Send the complete response
|
||||
if final_text:
|
||||
yield final_text
|
||||
else:
|
||||
logger.warning("No text in final result")
|
||||
|
||||
# Log completion metrics
|
||||
logger.info(
|
||||
"Chatbot response completed",
|
||||
total_tokens=result.usage().total_tokens if hasattr(result, "usage") else None,
|
||||
)
|
||||
|
||||
except Exception as e:
|
||||
logger.error("Chatbot streaming failed", error=str(e))
|
||||
yield f"\n\n[Error: {str(e)}]"
|
||||
|
||||
|
||||
# Global chatbot agent instance
|
||||
_chatbot_agent: ChatbotAgent | None = None
|
||||
|
||||
|
||||
def get_chatbot_agent() -> ChatbotAgent:
|
||||
"""Get or create the global chatbot agent instance."""
|
||||
global _chatbot_agent
|
||||
if _chatbot_agent is None:
|
||||
_chatbot_agent = ChatbotAgent()
|
||||
return _chatbot_agent
|
||||
|
||||
|
||||
async def chatbot_stream(request: ChatbotRequest) -> StreamingResponse:
|
||||
"""
|
||||
FastAPI endpoint handler for streaming chatbot responses.
|
||||
|
||||
Args:
|
||||
request: ChatbotRequest with message and settings
|
||||
|
||||
Returns:
|
||||
StreamingResponse with text/event-stream content
|
||||
"""
|
||||
agent = get_chatbot_agent()
|
||||
|
||||
# Ensure MCP HTTP server is running
|
||||
await agent.start_mcp_server()
|
||||
|
||||
# Stream the response (connects to MCP server via HTTP)
|
||||
return StreamingResponse(
|
||||
agent.stream_chat_response(request),
|
||||
media_type="text/plain",
|
||||
)
|
||||
@@ -43,6 +43,7 @@ pillow = "^11.3.0"
|
||||
opentelemetry-api = "^1.38.0"
|
||||
opentelemetry-sdk = "^1.38.0"
|
||||
opentelemetry-exporter-otlp-proto-grpc = "^1.38.0"
|
||||
presidio-analyzer = "^2.2.360"
|
||||
|
||||
[tool.poetry.group.dev.dependencies]
|
||||
ruff = "^0.9.2"
|
||||
|
||||
@@ -12,6 +12,7 @@ import {
|
||||
HelpCircle,
|
||||
Key,
|
||||
LayoutDashboard,
|
||||
MessageSquare,
|
||||
Search,
|
||||
Settings,
|
||||
Siren,
|
||||
@@ -41,6 +42,7 @@ import ThemeToggle from './components/ThemeToggle';
|
||||
import YaraRulesManager from './components/Yara/YaraManager';
|
||||
import Containers from './components/Containers/Containers';
|
||||
import AgentsPage from './components/Agents/AgentsPage';
|
||||
import ChatbotPage from './components/Chatbot/ChatbotPage';
|
||||
import FileBrowser from './components/FileBrowser/FileBrowser';
|
||||
import Chromium from './components/Chromium/Chromium';
|
||||
import Dpapi from './components/Dpapi/Dpapi';
|
||||
@@ -197,9 +199,13 @@ const Sidebar = ({ onCollapse }) => {
|
||||
{ id: 'reporting', label: 'Reporting', icon: BarChart2, path: '/reporting' }
|
||||
];
|
||||
|
||||
// Add Agents tab if LiteLLM is available
|
||||
// Add Chatbot and Agents tabs if LiteLLM is available
|
||||
const navigationItems = litellmAvailable
|
||||
? [...baseNavigationItems, { id: 'agents', label: 'Agents', icon: Bot, path: '/agents' }]
|
||||
? [
|
||||
...baseNavigationItems,
|
||||
{ id: 'chatbot', label: 'Chatbot', icon: MessageSquare, path: '/chatbot' },
|
||||
{ id: 'agents', label: 'Agents', icon: Bot, path: '/agents' }
|
||||
]
|
||||
: baseNavigationItems;
|
||||
|
||||
const utilityItems = [
|
||||
@@ -369,6 +375,7 @@ const App = () => {
|
||||
<Route path="/dpapi" element={<Dpapi />} />
|
||||
<Route path="/yara-rules" element={<YaraRulesManager />} />
|
||||
<Route path="/containers" element={<Containers />} />
|
||||
<Route path="/chatbot" element={<ChatbotPage />} />
|
||||
<Route path="/agents" element={<AgentsPage />} />
|
||||
<Route path="/reporting" element={<ReportingPage />} />
|
||||
<Route path="/reporting/source/:sourceName" element={<SourceReportPage />} />
|
||||
|
||||
@@ -0,0 +1,520 @@
|
||||
import { AlertCircle, Bot, Send, Settings as SettingsIcon, Trash2, X } from 'lucide-react';
|
||||
import { useEffect, useRef, useState } from 'react';
|
||||
import ExampleQueries from './ExampleQueries';
|
||||
import MessageBubble from './MessageBubble';
|
||||
|
||||
const MESSAGES_STORAGE_KEY = 'chatbot_messages';
|
||||
const TOKEN_WARNING_THRESHOLD = 50000;
|
||||
|
||||
// Rough token estimation: ~4 characters per token
|
||||
const estimateTokens = (text) => {
|
||||
return Math.ceil(text.length / 4);
|
||||
};
|
||||
|
||||
const ChatbotPage = () => {
|
||||
const [messages, setMessages] = useState(() => {
|
||||
// Load messages from localStorage on mount
|
||||
try {
|
||||
const saved = localStorage.getItem(MESSAGES_STORAGE_KEY);
|
||||
return saved ? JSON.parse(saved) : [];
|
||||
} catch {
|
||||
return [];
|
||||
}
|
||||
});
|
||||
const [currentMessage, setCurrentMessage] = useState('');
|
||||
const [isStreaming, setIsStreaming] = useState(false);
|
||||
const [useHistory, setUseHistory] = useState(true);
|
||||
const [temperature, setTemperature] = useState(0.7);
|
||||
const [error, setError] = useState(null);
|
||||
const [showSettings, setShowSettings] = useState(false);
|
||||
const [systemPrompt, setSystemPrompt] = useState('');
|
||||
const [originalPrompt, setOriginalPrompt] = useState('');
|
||||
const [savingPrompt, setSavingPrompt] = useState(false);
|
||||
const [promptError, setPromptError] = useState(null);
|
||||
const [spendData, setSpendData] = useState(null);
|
||||
const [showTokenWarning, setShowTokenWarning] = useState(false);
|
||||
const [estimatedTokens, setEstimatedTokens] = useState(0);
|
||||
|
||||
const messagesEndRef = useRef(null);
|
||||
const abortControllerRef = useRef(null);
|
||||
const textareaRef = useRef(null);
|
||||
|
||||
// Auto-scroll to bottom when messages change
|
||||
const scrollToBottom = () => {
|
||||
messagesEndRef.current?.scrollIntoView({ behavior: 'smooth' });
|
||||
};
|
||||
|
||||
useEffect(() => {
|
||||
scrollToBottom();
|
||||
}, [messages]);
|
||||
|
||||
// Save messages to localStorage whenever they change
|
||||
useEffect(() => {
|
||||
try {
|
||||
localStorage.setItem(MESSAGES_STORAGE_KEY, JSON.stringify(messages));
|
||||
} catch (err) {
|
||||
console.error('Failed to save messages to localStorage:', err);
|
||||
}
|
||||
}, [messages]);
|
||||
|
||||
// Refocus input after response completes
|
||||
useEffect(() => {
|
||||
if (!isStreaming && messages.length > 0) {
|
||||
// Small delay to ensure DOM has updated
|
||||
setTimeout(() => {
|
||||
textareaRef.current?.focus();
|
||||
}, 100);
|
||||
}
|
||||
}, [isStreaming, messages.length]);
|
||||
|
||||
// Fetch system prompt when settings are opened
|
||||
useEffect(() => {
|
||||
if (showSettings && !systemPrompt) {
|
||||
fetchSystemPrompt();
|
||||
}
|
||||
}, [showSettings]);
|
||||
|
||||
// Fetch spend data on mount
|
||||
useEffect(() => {
|
||||
fetchSpendData();
|
||||
}, []);
|
||||
|
||||
// Calculate token usage and show warning if needed
|
||||
useEffect(() => {
|
||||
const totalText = messages.map(m => m.content).join(' ');
|
||||
const tokens = estimateTokens(totalText);
|
||||
setEstimatedTokens(tokens);
|
||||
|
||||
if (tokens > TOKEN_WARNING_THRESHOLD && !showTokenWarning) {
|
||||
setShowTokenWarning(true);
|
||||
}
|
||||
}, [messages]);
|
||||
|
||||
const fetchSpendData = async () => {
|
||||
try {
|
||||
const response = await fetch('/api/agents/spend-data');
|
||||
if (!response.ok) throw new Error('Network response error');
|
||||
const result = await response.json();
|
||||
|
||||
setSpendData({
|
||||
spend: result.total_spend,
|
||||
total_tokens: result.total_tokens
|
||||
});
|
||||
} catch (err) {
|
||||
console.error('Error fetching spend data:', err);
|
||||
// Don't set error state - gracefully degrade
|
||||
}
|
||||
};
|
||||
|
||||
const fetchSystemPrompt = async () => {
|
||||
try {
|
||||
const query = {
|
||||
query: `
|
||||
query GetChatbotPrompt {
|
||||
agent_prompts_by_pk(name: "chatbot") {
|
||||
name
|
||||
prompt
|
||||
description
|
||||
}
|
||||
}
|
||||
`
|
||||
};
|
||||
|
||||
const response = await fetch('/hasura/v1/graphql', {
|
||||
method: 'POST',
|
||||
headers: {
|
||||
'Content-Type': 'application/json',
|
||||
'x-hasura-admin-secret': window.ENV.HASURA_ADMIN_SECRET,
|
||||
},
|
||||
body: JSON.stringify(query)
|
||||
});
|
||||
|
||||
if (!response.ok) throw new Error('Network response error');
|
||||
const result = await response.json();
|
||||
if (result.errors) throw new Error(result.errors[0].message);
|
||||
|
||||
const prompt = result.data.agent_prompts_by_pk?.prompt || '';
|
||||
setSystemPrompt(prompt);
|
||||
setOriginalPrompt(prompt);
|
||||
} catch (err) {
|
||||
console.error('Error fetching system prompt:', err);
|
||||
setPromptError('Failed to load system prompt');
|
||||
}
|
||||
};
|
||||
|
||||
const saveSystemPrompt = async () => {
|
||||
if (systemPrompt === originalPrompt) return;
|
||||
|
||||
setSavingPrompt(true);
|
||||
setPromptError(null);
|
||||
|
||||
try {
|
||||
const mutation = {
|
||||
query: `
|
||||
mutation UpsertChatbotPrompt($prompt: String!) {
|
||||
insert_agent_prompts_one(
|
||||
object: {
|
||||
name: "chatbot",
|
||||
prompt: $prompt,
|
||||
description: "Interactive chatbot for querying Nemesis data"
|
||||
},
|
||||
on_conflict: {
|
||||
constraint: agent_prompts_pkey,
|
||||
update_columns: [prompt]
|
||||
}
|
||||
) {
|
||||
name
|
||||
prompt
|
||||
}
|
||||
}
|
||||
`,
|
||||
variables: { prompt: systemPrompt }
|
||||
};
|
||||
|
||||
const response = await fetch('/hasura/v1/graphql', {
|
||||
method: 'POST',
|
||||
headers: {
|
||||
'Content-Type': 'application/json',
|
||||
'x-hasura-admin-secret': window.ENV.HASURA_ADMIN_SECRET,
|
||||
},
|
||||
body: JSON.stringify(mutation)
|
||||
});
|
||||
|
||||
if (!response.ok) throw new Error('Network response error');
|
||||
const result = await response.json();
|
||||
if (result.errors) throw new Error(result.errors[0].message);
|
||||
|
||||
setOriginalPrompt(systemPrompt);
|
||||
alert('System prompt saved successfully!');
|
||||
} catch (err) {
|
||||
console.error('Error saving system prompt:', err);
|
||||
setPromptError('Failed to save system prompt');
|
||||
} finally {
|
||||
setSavingPrompt(false);
|
||||
}
|
||||
};
|
||||
|
||||
const cancelPromptEdit = () => {
|
||||
setSystemPrompt(originalPrompt);
|
||||
setPromptError(null);
|
||||
};
|
||||
|
||||
const sendMessage = async (messageText) => {
|
||||
if (!messageText.trim() || isStreaming) return;
|
||||
|
||||
const userMessage = { role: 'user', content: messageText.trim() };
|
||||
setMessages(prev => [...prev, userMessage]);
|
||||
setCurrentMessage('');
|
||||
setIsStreaming(true);
|
||||
setError(null);
|
||||
|
||||
// Create abort controller for this request
|
||||
abortControllerRef.current = new AbortController();
|
||||
|
||||
try {
|
||||
const response = await fetch('/api/chatbot/stream', {
|
||||
method: 'POST',
|
||||
headers: { 'Content-Type': 'application/json' },
|
||||
body: JSON.stringify({
|
||||
message: messageText.trim(),
|
||||
history: useHistory ? messages : [],
|
||||
use_history: useHistory,
|
||||
temperature: temperature
|
||||
}),
|
||||
signal: abortControllerRef.current.signal
|
||||
});
|
||||
|
||||
if (!response.ok) {
|
||||
throw new Error(`HTTP ${response.status}: ${response.statusText}`);
|
||||
}
|
||||
|
||||
const reader = response.body.getReader();
|
||||
const decoder = new TextDecoder();
|
||||
let assistantMessage = '';
|
||||
|
||||
// Add empty assistant message that we'll update
|
||||
setMessages(prev => [...prev, { role: 'assistant', content: '' }]);
|
||||
|
||||
while (true) {
|
||||
const { done, value } = await reader.read();
|
||||
if (done) break;
|
||||
|
||||
const chunk = decoder.decode(value, { stream: true });
|
||||
assistantMessage += chunk;
|
||||
|
||||
// Update the last message (assistant's response) in real-time
|
||||
setMessages(prev => {
|
||||
const updated = [...prev];
|
||||
updated[updated.length - 1] = {
|
||||
role: 'assistant',
|
||||
content: assistantMessage
|
||||
};
|
||||
return updated;
|
||||
});
|
||||
}
|
||||
|
||||
} catch (err) {
|
||||
if (err.name === 'AbortError') {
|
||||
console.log('Request aborted');
|
||||
} else {
|
||||
console.error('Streaming error:', err);
|
||||
setError(err.message || 'Failed to get response from chatbot');
|
||||
|
||||
// Remove the empty assistant message if there was an error
|
||||
setMessages(prev => prev.filter(msg => msg.content !== ''));
|
||||
}
|
||||
} finally {
|
||||
setIsStreaming(false);
|
||||
abortControllerRef.current = null;
|
||||
// Refresh spend data after response completes
|
||||
fetchSpendData();
|
||||
}
|
||||
};
|
||||
|
||||
const handleSubmit = (e) => {
|
||||
e.preventDefault();
|
||||
sendMessage(currentMessage);
|
||||
};
|
||||
|
||||
const handleExampleClick = (example) => {
|
||||
setCurrentMessage(example);
|
||||
};
|
||||
|
||||
const clearHistory = () => {
|
||||
if (window.confirm('Clear all conversation history?')) {
|
||||
setMessages([]);
|
||||
setError(null);
|
||||
setShowTokenWarning(false);
|
||||
// Clear from localStorage as well
|
||||
try {
|
||||
localStorage.removeItem(MESSAGES_STORAGE_KEY);
|
||||
} catch (err) {
|
||||
console.error('Failed to clear messages from localStorage:', err);
|
||||
}
|
||||
}
|
||||
};
|
||||
|
||||
const handleKeyPress = (e) => {
|
||||
if (e.key === 'Enter' && !e.shiftKey) {
|
||||
e.preventDefault();
|
||||
handleSubmit(e);
|
||||
}
|
||||
};
|
||||
|
||||
return (
|
||||
<div className="flex flex-col h-[calc(100vh-1rem)] max-w-6xl mx-auto px-4 pt-4 pb-2">
|
||||
{/* Header */}
|
||||
<div className="flex items-center justify-between mb-4 pb-4 border-b border-gray-200 dark:border-gray-700">
|
||||
<div className="flex items-center space-x-3">
|
||||
<Bot className="w-8 h-8 text-blue-500" />
|
||||
<div>
|
||||
<h1 className="text-2xl font-bold text-gray-800 dark:text-white">Chatbot</h1>
|
||||
<p className="text-sm text-gray-600 dark:text-gray-400">
|
||||
Query Nemesis data with natural language
|
||||
</p>
|
||||
</div>
|
||||
</div>
|
||||
|
||||
{/* LLM Usage Stats */}
|
||||
{spendData && (
|
||||
<div className="flex items-center space-x-4 text-sm text-white">
|
||||
<div className="text-center">
|
||||
<div className="font-semibold">
|
||||
${spendData.spend ? spendData.spend.toFixed(4) : '0.0000'}
|
||||
</div>
|
||||
<div className="text-xs opacity-80">Total Spend</div>
|
||||
</div>
|
||||
<div className="text-center">
|
||||
<div className="font-semibold">
|
||||
{spendData.total_tokens ? spendData.total_tokens.toLocaleString() : '0'}
|
||||
</div>
|
||||
<div className="text-xs opacity-80">Total Tokens</div>
|
||||
</div>
|
||||
</div>
|
||||
)}
|
||||
|
||||
<div className="flex items-center space-x-3">
|
||||
{/* Settings Toggle */}
|
||||
<button
|
||||
onClick={() => setShowSettings(!showSettings)}
|
||||
className={`p-2 rounded-lg transition-colors ${showSettings
|
||||
? 'bg-blue-100 dark:bg-blue-900 text-blue-600 dark:text-blue-300'
|
||||
: 'hover:bg-gray-100 dark:hover:bg-gray-700 text-gray-600 dark:text-gray-400'
|
||||
}`}
|
||||
title="Settings"
|
||||
>
|
||||
<SettingsIcon className="w-5 h-5" />
|
||||
</button>
|
||||
|
||||
{/* Clear History */}
|
||||
<button
|
||||
onClick={clearHistory}
|
||||
disabled={messages.length === 0}
|
||||
className="flex items-center space-x-1 px-3 py-2 text-sm bg-red-600 text-white rounded-lg hover:bg-red-700 disabled:bg-gray-400 disabled:cursor-not-allowed transition-colors"
|
||||
title="Clear conversation history"
|
||||
>
|
||||
<Trash2 className="w-4 h-4" />
|
||||
<span>Clear History</span>
|
||||
</button>
|
||||
</div>
|
||||
</div>
|
||||
|
||||
{/* Settings Panel */}
|
||||
{showSettings && (
|
||||
<div className="mb-4 p-4 bg-gray-50 dark:bg-gray-800 rounded-lg border border-gray-200 dark:border-gray-700 max-h-96 overflow-y-auto flex-shrink-0">
|
||||
<h3 className="text-sm font-semibold text-gray-700 dark:text-gray-300 mb-3">Settings</h3>
|
||||
|
||||
<div className="space-y-4">
|
||||
{/* Temperature Control */}
|
||||
<div>
|
||||
<label className="flex justify-between text-sm text-gray-600 dark:text-gray-400 mb-1">
|
||||
<span>Temperature: {temperature.toFixed(1)}</span>
|
||||
<span className="text-xs text-gray-500">
|
||||
{temperature < 0.3 ? 'Focused' : temperature < 0.7 ? 'Balanced' : 'Creative'}
|
||||
</span>
|
||||
</label>
|
||||
<input
|
||||
type="range"
|
||||
min="0"
|
||||
max="1"
|
||||
step="0.1"
|
||||
value={temperature}
|
||||
onChange={(e) => setTemperature(parseFloat(e.target.value))}
|
||||
className="w-full h-2 bg-gray-200 rounded-lg appearance-none cursor-pointer dark:bg-gray-700"
|
||||
/>
|
||||
</div>
|
||||
|
||||
{/* Use History Toggle */}
|
||||
<div className="flex items-center justify-between">
|
||||
<span className="text-sm text-gray-600 dark:text-gray-400">Use Conversation History</span>
|
||||
<button
|
||||
onClick={() => setUseHistory(!useHistory)}
|
||||
className={`relative inline-flex h-6 w-11 items-center rounded-full transition-colors ${useHistory ? 'bg-blue-600' : 'bg-gray-300 dark:bg-gray-600'
|
||||
}`}
|
||||
>
|
||||
<span
|
||||
className={`inline-block h-4 w-4 transform rounded-full bg-white transition-transform ${useHistory ? 'translate-x-6' : 'translate-x-1'
|
||||
}`}
|
||||
/>
|
||||
</button>
|
||||
</div>
|
||||
|
||||
{/* System Prompt Editor */}
|
||||
<div className="pt-4 border-t border-gray-200 dark:border-gray-700">
|
||||
<h4 className="text-sm font-semibold text-gray-700 dark:text-gray-300 mb-2">System Prompt</h4>
|
||||
|
||||
{promptError && (
|
||||
<div className="mb-3 p-2 bg-red-50 dark:bg-red-900/20 border border-red-200 dark:border-red-800 rounded text-sm text-red-600 dark:text-red-400">
|
||||
{promptError}
|
||||
</div>
|
||||
)}
|
||||
|
||||
<textarea
|
||||
value={systemPrompt}
|
||||
onChange={(e) => setSystemPrompt(e.target.value)}
|
||||
className="w-full h-64 px-3 py-2 text-sm border border-gray-300 dark:border-gray-600 rounded-lg resize-none focus:ring-2 focus:ring-blue-500 focus:border-transparent bg-white dark:bg-gray-700 text-gray-900 dark:text-white font-mono"
|
||||
placeholder="Loading system prompt..."
|
||||
/>
|
||||
|
||||
<div className="flex space-x-2 mt-3">
|
||||
<button
|
||||
onClick={saveSystemPrompt}
|
||||
disabled={savingPrompt || systemPrompt === originalPrompt}
|
||||
className="px-4 py-2 text-sm bg-blue-600 text-white rounded-lg hover:bg-blue-700 disabled:bg-gray-400 disabled:cursor-not-allowed transition-colors"
|
||||
>
|
||||
{savingPrompt ? 'Saving...' : 'Save Prompt'}
|
||||
</button>
|
||||
<button
|
||||
onClick={cancelPromptEdit}
|
||||
disabled={savingPrompt || systemPrompt === originalPrompt}
|
||||
className="px-4 py-2 text-sm bg-gray-200 dark:bg-gray-700 text-gray-700 dark:text-gray-300 rounded-lg hover:bg-gray-300 dark:hover:bg-gray-600 disabled:opacity-50 disabled:cursor-not-allowed transition-colors"
|
||||
>
|
||||
Cancel
|
||||
</button>
|
||||
</div>
|
||||
</div>
|
||||
</div>
|
||||
</div>
|
||||
)}
|
||||
|
||||
{/* Scrollable Content Area */}
|
||||
<div className="flex-1 overflow-y-auto space-y-4 min-h-0">
|
||||
{/* Token Warning */}
|
||||
{showTokenWarning && (
|
||||
<div className="p-3 bg-yellow-50 dark:bg-yellow-900/20 border border-yellow-200 dark:border-yellow-800 rounded-lg flex items-start space-x-2">
|
||||
<AlertCircle className="w-5 h-5 text-yellow-600 dark:text-yellow-400 flex-shrink-0 mt-0.5" />
|
||||
<div className="flex-1">
|
||||
<p className="text-sm text-yellow-600 dark:text-yellow-400">
|
||||
<strong>High token usage detected:</strong> Your conversation history is getting long (~{estimatedTokens.toLocaleString()} tokens)
|
||||
and consuming significant LLM resources. Consider clicking "Clear History" to reset and reduce costs.
|
||||
</p>
|
||||
</div>
|
||||
<button
|
||||
onClick={() => setShowTokenWarning(false)}
|
||||
className="text-yellow-600 dark:text-yellow-400 hover:text-yellow-800 dark:hover:text-yellow-200"
|
||||
title="Dismiss warning"
|
||||
>
|
||||
<X className="w-5 h-5" />
|
||||
</button>
|
||||
</div>
|
||||
)}
|
||||
|
||||
{/* Error Display */}
|
||||
{error && (
|
||||
<div className="p-3 bg-red-50 dark:bg-red-900/20 border border-red-200 dark:border-red-800 rounded-lg flex items-start space-x-2">
|
||||
<AlertCircle className="w-5 h-5 text-red-600 dark:text-red-400 flex-shrink-0 mt-0.5" />
|
||||
<div className="flex-1">
|
||||
<p className="text-sm text-red-600 dark:text-red-400">{error}</p>
|
||||
</div>
|
||||
</div>
|
||||
)}
|
||||
|
||||
{/* Example Queries (show when no messages) */}
|
||||
{messages.length === 0 && (
|
||||
<ExampleQueries onExampleClick={handleExampleClick} />
|
||||
)}
|
||||
|
||||
{/* Messages */}
|
||||
<div className="space-y-4 px-2">
|
||||
{messages.map((msg, idx) => (
|
||||
<MessageBubble
|
||||
key={idx}
|
||||
message={msg}
|
||||
isStreaming={isStreaming && idx === messages.length - 1}
|
||||
/>
|
||||
))}
|
||||
<div ref={messagesEndRef} />
|
||||
</div>
|
||||
</div>
|
||||
|
||||
{/* Input Form */}
|
||||
<form onSubmit={handleSubmit} className="flex-shrink-0 border-t border-gray-200 dark:border-gray-700 pt-4 mt-4 pb-4">
|
||||
<div className="flex space-x-2">
|
||||
<textarea
|
||||
ref={textareaRef}
|
||||
value={currentMessage}
|
||||
onChange={(e) => setCurrentMessage(e.target.value)}
|
||||
onKeyPress={handleKeyPress}
|
||||
placeholder="Ask a question about your data..."
|
||||
disabled={isStreaming}
|
||||
className="flex-1 px-4 py-3 border border-gray-300 dark:border-gray-600 rounded-lg resize-none focus:ring-2 focus:ring-blue-500 focus:border-transparent bg-white dark:bg-gray-700 text-gray-900 dark:text-white placeholder-gray-500 dark:placeholder-gray-400 disabled:bg-gray-100 dark:disabled:bg-gray-800 disabled:cursor-not-allowed"
|
||||
rows={2}
|
||||
/>
|
||||
<button
|
||||
type="submit"
|
||||
disabled={!currentMessage.trim() || isStreaming}
|
||||
className="px-6 py-3 bg-blue-600 text-white rounded-lg hover:bg-blue-700 disabled:bg-gray-400 disabled:cursor-not-allowed transition-colors flex items-center space-x-2"
|
||||
>
|
||||
<Send className="w-5 h-5" />
|
||||
<span>{isStreaming ? 'Sending...' : 'Send'}</span>
|
||||
</button>
|
||||
</div>
|
||||
<p className="mt-2 text-xs text-gray-500 dark:text-gray-400">
|
||||
Press Enter to send, Shift+Enter for new line
|
||||
</p>
|
||||
</form>
|
||||
</div>
|
||||
);
|
||||
};
|
||||
|
||||
export default ChatbotPage;
|
||||
@@ -0,0 +1,52 @@
|
||||
import { Lightbulb } from 'lucide-react';
|
||||
|
||||
const ExampleQueries = ({ onExampleClick }) => {
|
||||
const examples = [
|
||||
{
|
||||
title: 'High Severity Findings',
|
||||
query: 'How many findings have severity greater than 7?',
|
||||
description: 'Count critical security findings',
|
||||
},
|
||||
{
|
||||
title: 'Exposed Credentials',
|
||||
query: 'What decrypted passwords are available from host://WORKSTATION01?',
|
||||
description: 'Find credentials from a specific source',
|
||||
},
|
||||
{
|
||||
title: 'Lateral Movement',
|
||||
query: 'How can I access example.com?',
|
||||
description: 'Identify lateral movement opportunities',
|
||||
},
|
||||
];
|
||||
|
||||
return (
|
||||
<div className="mb-6 space-y-4">
|
||||
<div className="flex items-center space-x-2 text-gray-600 dark:text-gray-400">
|
||||
<Lightbulb className="w-5 h-5" />
|
||||
<h3 className="text-sm font-semibold">Try these example queries:</h3>
|
||||
</div>
|
||||
|
||||
<div className="grid grid-cols-1 md:grid-cols-3 gap-4">
|
||||
{examples.map((example, idx) => (
|
||||
<button
|
||||
key={idx}
|
||||
onClick={() => onExampleClick(example.query)}
|
||||
className="text-left p-4 bg-white dark:bg-gray-800 border border-gray-200 dark:border-gray-700 rounded-lg hover:border-blue-500 dark:hover:border-blue-400 hover:shadow-md transition-all group"
|
||||
>
|
||||
<h4 className="font-semibold text-sm text-gray-800 dark:text-gray-200 mb-1 group-hover:text-blue-600 dark:group-hover:text-blue-400">
|
||||
{example.title}
|
||||
</h4>
|
||||
<p className="text-xs text-gray-600 dark:text-gray-400 mb-2">
|
||||
{example.description}
|
||||
</p>
|
||||
<p className="text-xs text-gray-500 dark:text-gray-500 italic line-clamp-2">
|
||||
"{example.query}"
|
||||
</p>
|
||||
</button>
|
||||
))}
|
||||
</div>
|
||||
</div>
|
||||
);
|
||||
};
|
||||
|
||||
export default ExampleQueries;
|
||||
@@ -0,0 +1,91 @@
|
||||
import React from 'react';
|
||||
import { Bot, User } from 'lucide-react';
|
||||
import ReactMarkdown from 'react-markdown';
|
||||
|
||||
const TypingIndicator = () => (
|
||||
<div className="flex space-x-1">
|
||||
<div className="w-2 h-2 bg-gray-400 rounded-full animate-bounce" style={{ animationDelay: '0ms' }}></div>
|
||||
<div className="w-2 h-2 bg-gray-400 rounded-full animate-bounce" style={{ animationDelay: '150ms' }}></div>
|
||||
<div className="w-2 h-2 bg-gray-400 rounded-full animate-bounce" style={{ animationDelay: '300ms' }}></div>
|
||||
</div>
|
||||
);
|
||||
|
||||
const MessageBubble = ({ message, isStreaming }) => {
|
||||
const isUser = message.role === 'user';
|
||||
|
||||
return (
|
||||
<div className={`flex ${isUser ? 'justify-end' : 'justify-start'}`}>
|
||||
<div className={`flex max-w-3xl ${isUser ? 'flex-row-reverse' : 'flex-row'} items-start space-x-3`}>
|
||||
{/* Avatar */}
|
||||
<div
|
||||
className={`flex-shrink-0 w-8 h-8 rounded-full flex items-center justify-center ${
|
||||
isUser
|
||||
? 'bg-blue-600 text-white'
|
||||
: 'bg-gray-200 dark:bg-gray-700 text-gray-600 dark:text-gray-300'
|
||||
}`}
|
||||
>
|
||||
{isUser ? <User className="w-5 h-5" /> : <Bot className="w-5 h-5" />}
|
||||
</div>
|
||||
|
||||
{/* Message Content */}
|
||||
<div
|
||||
className={`flex-1 px-4 py-3 rounded-lg ${
|
||||
isUser
|
||||
? 'bg-blue-600 text-white ml-3'
|
||||
: 'bg-gray-100 dark:bg-gray-800 text-gray-900 dark:text-gray-100 mr-3'
|
||||
}`}
|
||||
>
|
||||
{isUser ? (
|
||||
// User messages are plain text
|
||||
<p className="text-sm whitespace-pre-wrap break-words">{message.content}</p>
|
||||
) : (
|
||||
// Assistant messages support markdown
|
||||
<>
|
||||
{message.content ? (
|
||||
<div className="prose prose-sm dark:prose-invert max-w-none">
|
||||
<ReactMarkdown
|
||||
components={{
|
||||
// Customize code blocks
|
||||
code({ node, inline, className, children, ...props }) {
|
||||
return inline ? (
|
||||
<code
|
||||
className="px-1 py-0.5 bg-gray-200 dark:bg-gray-700 rounded text-xs"
|
||||
{...props}
|
||||
>
|
||||
{children}
|
||||
</code>
|
||||
) : (
|
||||
<pre className="bg-gray-200 dark:bg-gray-900 p-3 rounded-lg overflow-x-auto">
|
||||
<code className="text-xs" {...props}>
|
||||
{children}
|
||||
</code>
|
||||
</pre>
|
||||
);
|
||||
},
|
||||
// Customize tables
|
||||
table({ children }) {
|
||||
return (
|
||||
<div className="overflow-x-auto">
|
||||
<table className="min-w-full divide-y divide-gray-300 dark:divide-gray-600">
|
||||
{children}
|
||||
</table>
|
||||
</div>
|
||||
);
|
||||
},
|
||||
}}
|
||||
>
|
||||
{message.content}
|
||||
</ReactMarkdown>
|
||||
</div>
|
||||
) : isStreaming ? (
|
||||
<TypingIndicator />
|
||||
) : null}
|
||||
</>
|
||||
)}
|
||||
</div>
|
||||
</div>
|
||||
</div>
|
||||
);
|
||||
};
|
||||
|
||||
export default MessageBubble;
|
||||
@@ -0,0 +1,117 @@
|
||||
import React, { useState } from 'react';
|
||||
import { ChevronDown, ChevronRight, Copy, Check } from 'lucide-react';
|
||||
|
||||
const QueryModal = ({ queries }) => {
|
||||
const [expandedQueries, setExpandedQueries] = useState({});
|
||||
const [copiedIndex, setCopiedIndex] = useState(null);
|
||||
|
||||
const toggleQuery = (index) => {
|
||||
setExpandedQueries(prev => ({
|
||||
...prev,
|
||||
[index]: !prev[index]
|
||||
}));
|
||||
};
|
||||
|
||||
const copyToClipboard = async (text, index) => {
|
||||
try {
|
||||
await navigator.clipboard.writeText(text);
|
||||
setCopiedIndex(index);
|
||||
setTimeout(() => setCopiedIndex(null), 2000);
|
||||
} catch (err) {
|
||||
console.error('Failed to copy:', err);
|
||||
}
|
||||
};
|
||||
|
||||
if (!queries || queries.length === 0) return null;
|
||||
|
||||
return (
|
||||
<div className="mb-4 border border-gray-200 dark:border-gray-700 rounded-lg bg-white dark:bg-gray-800 shadow-sm">
|
||||
<div className="p-3 border-b border-gray-200 dark:border-gray-700">
|
||||
<h3 className="text-sm font-semibold text-gray-700 dark:text-gray-300">
|
||||
SQL Queries Executed
|
||||
</h3>
|
||||
</div>
|
||||
|
||||
<div className="divide-y divide-gray-200 dark:divide-gray-700">
|
||||
{queries.map((query, index) => {
|
||||
const isExpanded = expandedQueries[index];
|
||||
const isCopied = copiedIndex === index;
|
||||
|
||||
return (
|
||||
<div key={index} className="transition-colors">
|
||||
{/* Query Header - Collapsible */}
|
||||
<button
|
||||
onClick={() => toggleQuery(index)}
|
||||
className="w-full flex items-center justify-between p-3 hover:bg-gray-50 dark:hover:bg-gray-700/50 transition-colors"
|
||||
>
|
||||
<div className="flex items-center space-x-2 flex-1 min-w-0">
|
||||
{isExpanded ? (
|
||||
<ChevronDown className="w-4 h-4 text-gray-500 flex-shrink-0" />
|
||||
) : (
|
||||
<ChevronRight className="w-4 h-4 text-gray-500 flex-shrink-0" />
|
||||
)}
|
||||
<span className="text-sm font-medium text-gray-700 dark:text-gray-300 truncate">
|
||||
{query.name || `Query ${index + 1}`}
|
||||
</span>
|
||||
{query.timestamp && (
|
||||
<span className="text-xs text-gray-500 dark:text-gray-400">
|
||||
{new Date(query.timestamp).toLocaleTimeString()}
|
||||
</span>
|
||||
)}
|
||||
</div>
|
||||
{query.rowCount !== undefined && (
|
||||
<span className="text-xs text-gray-500 dark:text-gray-400 ml-2">
|
||||
{query.rowCount} rows
|
||||
</span>
|
||||
)}
|
||||
</button>
|
||||
|
||||
{/* Query Content - Expandable */}
|
||||
{isExpanded && (
|
||||
<div className="px-3 pb-3">
|
||||
<div className="relative">
|
||||
{/* Copy Button */}
|
||||
<button
|
||||
onClick={() => copyToClipboard(query.sql, index)}
|
||||
className="absolute top-2 right-2 p-1.5 bg-gray-100 dark:bg-gray-700 hover:bg-gray-200 dark:hover:bg-gray-600 rounded transition-colors"
|
||||
title="Copy to clipboard"
|
||||
>
|
||||
{isCopied ? (
|
||||
<Check className="w-4 h-4 text-green-600 dark:text-green-400" />
|
||||
) : (
|
||||
<Copy className="w-4 h-4 text-gray-600 dark:text-gray-400" />
|
||||
)}
|
||||
</button>
|
||||
|
||||
{/* SQL Code Block */}
|
||||
<pre className="bg-gray-50 dark:bg-gray-900 p-3 pr-12 rounded-lg overflow-x-auto text-xs font-mono">
|
||||
<code className="text-gray-800 dark:text-gray-200">
|
||||
{query.sql}
|
||||
</code>
|
||||
</pre>
|
||||
</div>
|
||||
|
||||
{/* Query Metadata */}
|
||||
{(query.executionTime || query.error) && (
|
||||
<div className="mt-2 text-xs text-gray-500 dark:text-gray-400">
|
||||
{query.executionTime && (
|
||||
<span>Execution time: {query.executionTime}ms</span>
|
||||
)}
|
||||
{query.error && (
|
||||
<span className="text-red-600 dark:text-red-400">
|
||||
Error: {query.error}
|
||||
</span>
|
||||
)}
|
||||
</div>
|
||||
)}
|
||||
</div>
|
||||
)}
|
||||
</div>
|
||||
);
|
||||
})}
|
||||
</div>
|
||||
</div>
|
||||
);
|
||||
};
|
||||
|
||||
export default QueryModal;
|
||||
Generated
+1
-1
@@ -3809,4 +3809,4 @@ propcache = ">=0.2.1"
|
||||
[metadata]
|
||||
lock-version = "2.1"
|
||||
python-versions = ">=3.12,<3.14"
|
||||
content-hash = "24ef2f823348ff25c13ffdc741e07ad735a7c6d9706faf37c69ba93843c21c42"
|
||||
content-hash = "f8e33433b5ee89307bdbea073cb80ce5b7c45eb4b872888479f10bf7e6ef6539"
|
||||
|
||||
@@ -26,6 +26,7 @@ watchdog = "^6.0.0"
|
||||
pyyaml = "^6.0.2"
|
||||
jinja2 = "^3.1.6"
|
||||
markdown = "^3.9"
|
||||
httpx = "^0.28.1"
|
||||
|
||||
[tool.poetry.group.dev.dependencies]
|
||||
pytest = "^8.4.2"
|
||||
|
||||
@@ -11,6 +11,7 @@ from pathlib import Path as PathLib
|
||||
from typing import Annotated
|
||||
from urllib.parse import urlparse
|
||||
|
||||
import httpx
|
||||
import psycopg
|
||||
import requests
|
||||
from common.db import get_postgres_connection_str
|
||||
@@ -45,7 +46,7 @@ from psycopg_pool import ConnectionPool
|
||||
from pydantic import ValidationError
|
||||
from web_api.container_monitor import get_monitor, start_monitor, stop_monitor
|
||||
from web_api.large_containers import LargeContainerProcessor
|
||||
from web_api.models.requests import CleanupRequest, EnrichmentRequest
|
||||
from web_api.models.requests import ChatbotRequest, CleanupRequest, EnrichmentRequest
|
||||
from web_api.models.responses import (
|
||||
ContainerStatusResponse,
|
||||
ContainerSubmissionResponse,
|
||||
@@ -1421,6 +1422,60 @@ async def run_translation(
|
||||
raise HTTPException(status_code=500, detail=str(e)) from e
|
||||
|
||||
|
||||
@app.post(
|
||||
"/chatbot/stream",
|
||||
tags=["chatbot"],
|
||||
summary="Stream chatbot responses",
|
||||
description="Stream interactive chatbot responses for querying Nemesis data",
|
||||
)
|
||||
async def chatbot_stream(
|
||||
request: ChatbotRequest = Body(..., description="Chatbot request with message and conversation history"),
|
||||
):
|
||||
"""
|
||||
Stream chatbot responses via Dapr to agents service.
|
||||
Uses HTTP streaming to provide real-time token-by-token responses.
|
||||
"""
|
||||
try:
|
||||
url = f"http://localhost:{DAPR_PORT}/v1.0/invoke/agents/method/agents/chatbot/stream"
|
||||
|
||||
logger.debug("Proxying chatbot request to agents service", message_length=len(request.message))
|
||||
|
||||
async def stream_proxy():
|
||||
"""Generator function that streams chunks from agents service."""
|
||||
try:
|
||||
async with httpx.AsyncClient() as client:
|
||||
async with client.stream(
|
||||
"POST",
|
||||
url,
|
||||
json=request.model_dump(),
|
||||
timeout=120.0,
|
||||
) as response:
|
||||
response.raise_for_status()
|
||||
async for chunk in response.aiter_bytes():
|
||||
if chunk:
|
||||
yield chunk
|
||||
except httpx.HTTPStatusError as e:
|
||||
logger.error("HTTP error from agents service", status_code=e.response.status_code)
|
||||
error_msg = f"\n\n[Error: Agents service returned {e.response.status_code}]"
|
||||
yield error_msg.encode()
|
||||
except httpx.TimeoutException:
|
||||
logger.error("Timeout streaming from agents service")
|
||||
yield b"\n\n[Error: Request timeout]"
|
||||
except Exception as e:
|
||||
logger.exception("Error in stream proxy")
|
||||
error_msg = f"\n\n[Error: {str(e)}]"
|
||||
yield error_msg.encode()
|
||||
|
||||
return StreamingResponse(
|
||||
stream_proxy(),
|
||||
media_type="text/plain",
|
||||
)
|
||||
|
||||
except Exception as e:
|
||||
logger.exception(message="Error initiating chatbot stream")
|
||||
raise HTTPException(status_code=500, detail=str(e)) from e
|
||||
|
||||
|
||||
@app.get(
|
||||
"/system/available-services",
|
||||
tags=["system"],
|
||||
|
||||
@@ -7,3 +7,19 @@ class EnrichmentRequest(BaseModel):
|
||||
|
||||
class CleanupRequest(BaseModel):
|
||||
expiration: str | None = None # ISO datetime or "all"
|
||||
|
||||
|
||||
class ChatbotMessage(BaseModel):
|
||||
"""A single chat message."""
|
||||
|
||||
role: str # "user" or "assistant"
|
||||
content: str
|
||||
|
||||
|
||||
class ChatbotRequest(BaseModel):
|
||||
"""Request model for chatbot queries."""
|
||||
|
||||
message: str
|
||||
history: list[ChatbotMessage] = []
|
||||
use_history: bool = True
|
||||
temperature: float = 0.7
|
||||
|
||||
Reference in New Issue
Block a user