diff --git a/.env.example b/.env.example index 9d68bb7ff..439c7027e 100644 --- a/.env.example +++ b/.env.example @@ -184,3 +184,21 @@ OWASP_WARMUP=1 # Remote MCP server: how many questions may be answered at the same time (each in its own thread). OPENCRE_MCP_ASK_CONCURRENCY=8 + +# --- Accounts and roles (web chatbot and MCP share one model; see docs/roles.md) --- +# Two roles: "user" = any account on an allowed domain; "admin" = the named accounts below. +# Leave ADMIN_EMAILS unset for the old behaviour (LOGIN_ALLOWED_EMAILS is the one allowlist, all admins). +# ADMIN_EMAILS=alice@owasp.org,bob@owasp.org +# LOGIN_ALLOWED_DOMAINS=owasp.org + +# --- Per-account question limits (chatbot + MCP ask; shared through Redis) --- +# Defaults per account: 5 / minute, 40 / hour, 150 / day; admins x4. Uncomment to change. +# RATE_LIMIT_PER_MINUTE=5 +# RATE_LIMIT_PER_HOUR=40 +# RATE_LIMIT_PER_DAY=150 +# RATE_LIMIT_ADMIN_MULTIPLIER=4 +# RATE_LIMIT_EXEMPT_EMAILS=evalbot@owasp.org +# RATE_LIMIT_ENABLED=0 # switch limits off + +# Warm each web worker in the background at start so the first chat is not slow (the MCP server always does). +# OWASP_WARMUP_WEB=1 diff --git a/application/__init__.py b/application/__init__.py index afa979655..1c318ee65 100644 --- a/application/__init__.py +++ b/application/__init__.py @@ -92,4 +92,8 @@ def create_app(mode: str = "production", conf: any = None) -> Any: compress.init_app(app) cache.init_app(app) + if mode != "test": + from application.utils.comprehension.shared import start_background_warm_up + + start_background_warm_up(app) # no-op unless OWASP_WARMUP_WEB=1 return app diff --git a/application/database/db.py b/application/database/db.py index 678bad10f..9bfe1b3f1 100644 --- a/application/database/db.py +++ b/application/database/db.py @@ -321,6 +321,29 @@ class AdminRuntimeSettingAudit(BaseModel): # type: ignore __table_args__ = (sqla.Index("ix_admin_runtime_setting_audit_key", "key"),) +class AnswerLog(BaseModel): # type: ignore + """Every answer the system gives, with an optional thumbs up/down from the asker.""" + + __tablename__ = "answer_log" + id = sqla.Column(sqla.String, primary_key=True, default=generate_uuid) + created_at = sqla.Column(sqla.DateTime, nullable=False) + source = sqla.Column(sqla.String, nullable=False) # "mcp" | "chat" + question = sqla.Column(sqla.Text, nullable=False) + answer = sqla.Column(sqla.Text, nullable=True) + citations = sqla.Column(sqla.Text, nullable=True) # JSON list of urls + channel = sqla.Column(sqla.String, nullable=True) + flags = sqla.Column(sqla.String, nullable=True) # active OWASP_CX_* flags + latency_ms = sqla.Column(sqla.Integer, nullable=True) + rating = sqla.Column(sqla.String, nullable=True) # "up" | "down" + comment = sqla.Column(sqla.Text, nullable=True) + rated_at = sqla.Column(sqla.DateTime, nullable=True) + + __table_args__ = ( + sqla.Index("ix_answer_log_created_at", "created_at"), + sqla.Index("ix_answer_log_rating", "rating"), + ) + + class User(BaseModel): # type: ignore """Persisted account identity for OIDC login (issue #586, RFC #876 TODO 1). diff --git a/application/frontend/src/components/DocumentNode/DocumentNode.tsx b/application/frontend/src/components/DocumentNode/DocumentNode.tsx index c8318e205..1326cc660 100644 --- a/application/frontend/src/components/DocumentNode/DocumentNode.tsx +++ b/application/frontend/src/components/DocumentNode/DocumentNode.tsx @@ -36,7 +36,8 @@ export interface DocumentNode { const linkTypesToNest = [TYPE_IS_PART_OF, TYPE_RELATED, TYPE_AUTOLINKED_TO]; const linkTypesExcludedInNesting = [TYPE_CONTAINS]; const linkTypesExcludedWhenNestingRelatedTo = [TYPE_RELATED, TYPE_IS_PART_OF, TYPE_CONTAINS]; -const linkDisplayOrder = [TYPE_LINKED_TO, TYPE_AUTOLINKED_TO, TYPE_CONTAINS, TYPE_IS_PART_OF, TYPE_RELATED]; +// "Related" is more useful than the hierarchy, so it comes before "child of". +const linkDisplayOrder = [TYPE_LINKED_TO, TYPE_AUTOLINKED_TO, TYPE_CONTAINS, TYPE_RELATED, TYPE_IS_PART_OF]; export const DocumentNode: FunctionComponent = ({ node, diff --git a/application/frontend/src/pages/Admin/Admin.test.tsx b/application/frontend/src/pages/Admin/Admin.test.tsx index 6cc33ae32..0ce005d03 100644 --- a/application/frontend/src/pages/Admin/Admin.test.tsx +++ b/application/frontend/src/pages/Admin/Admin.test.tsx @@ -122,6 +122,38 @@ describe('Admin', () => { expect(JSON.parse(put!.init.body)).toEqual({ interval_minutes: 30, enabled: true }); }); + it('edits project authority weights and saves only the changed ones', async () => { + loggedIn(); + const calls: { url: string; init?: any }[] = []; + (global as any).fetch = jest.fn((url: string, init?: any) => { + calls.push({ url: String(url), init }); + if (String(url).endsWith('/admin/authority') && init?.method === 'PUT') { + return jsonRes({ weights: { 'www-project-ai-exchange': 1.2 } }); + } + if (String(url).endsWith('/admin/authority')) { + return jsonRes({ + projects: [ + { slug: 'www-project-ai-exchange', name: 'ai exchange', weight: 1 }, + { slug: 'AISVS', name: 'aisvs', weight: 1 }, + ], + }); + } + return jsonRes({}); + }); + const { getByText, findByLabelText, findByText } = render( + + + + ); + fireEvent.click(getByText('Authority')); + const input = (await findByLabelText('www-project-ai-exchange weight')) as HTMLInputElement; + fireEvent.change(input, { target: { value: '1.2' } }); + fireEvent.click(getByText(/^Save/)); + expect(await findByText(/1 project weight\(s\) saved/)).toBeTruthy(); + const put = calls.find((c) => c.init?.method === 'PUT'); + expect(JSON.parse(put!.init.body)).toEqual({ weights: { 'www-project-ai-exchange': 1.2 } }); + }); + it('loads import runs and links to MyOpenCRE', async () => { loggedIn(); (global as any).fetch = jest.fn((url: string) => { diff --git a/application/frontend/src/pages/Admin/Admin.tsx b/application/frontend/src/pages/Admin/Admin.tsx index 2380f196a..48c65ca45 100644 --- a/application/frontend/src/pages/Admin/Admin.tsx +++ b/application/frontend/src/pages/Admin/Admin.tsx @@ -13,6 +13,7 @@ type Tab = | 'pipeline' | 'oie' | 'schedule' + | 'authority' | 'graph' | 'workbench' | 'myopencre' @@ -23,6 +24,10 @@ const TAB_META: Record = { config: { label: 'Config values', sub: 'Process env vs .env — admin only.' }, pipeline: { label: 'Pipeline', sub: 'Jobs, stage logs, ingest.' }, oie: { label: 'OIE queue', sub: 'RQ / OIE worker status.' }, + authority: { + label: 'Authority', + sub: 'Feature or soften projects. 1 = neutral; above 1 = always list the closest passage among the references; below 1 = rank lower; 0 = hidden.', + }, schedule: { label: 'Schedules', sub: 'How often the ingestion jobs run — change or switch off.' }, graph: { label: 'Review', sub: 'Staged links — expand section text, accept/reject.' }, workbench: { label: 'Workbench', sub: 'Accepted targets and re-ingest.' }, @@ -182,6 +187,7 @@ export const Admin = () => { { id: 'pipeline', label: 'Pipeline', group: 'Ops' }, { id: 'oie', label: 'OIE queue', group: 'Ops' }, { id: 'schedule', label: 'Schedules', group: 'Ops' }, + { id: 'authority', label: 'Authority', group: 'Ops' }, { id: 'graph', label: 'Review', group: 'Ops' }, { id: 'workbench', label: 'Workbench', group: 'Ops' }, { id: 'myopencre', label: 'MyOpenCRE', group: 'Apps' }, @@ -234,6 +240,7 @@ export const Admin = () => { {tab === 'pipeline' && } {tab === 'oie' && } {tab === 'schedule' && } + {tab === 'authority' && } {tab === 'graph' && } {tab === 'workbench' && } {tab === 'myopencre' && ( @@ -1074,6 +1081,109 @@ type ScheduleJob = { last_run: { status?: string; started_at?: string } | null; }; +type AuthorityRow = { slug: string; name: string; weight: number }; + +function AuthorityTab({ origin }: { origin: string }) { + const [rows, setRows] = useState([]); + const [draft, setDraft] = useState>({}); + const [filter, setFilter] = useState(''); + const [error, setError] = useState(null); + const [notice, setNotice] = useState(null); + const [busy, setBusy] = useState(false); + + const load = useCallback(() => { + fetch(`${origin}/admin/authority`) + .then(async (res) => { + const body = await readJson(res); + if (!res.ok) throw new Error(apiError(body, res)); + const list: AuthorityRow[] = body.projects || []; + setRows(list); + setDraft(Object.fromEntries(list.map((r) => [r.slug, String(r.weight)]))); + setError(null); + }) + .catch((err) => setError(String(err))); + }, [origin]); + + useEffect(() => { + load(); + }, [load]); + + const changed = rows.filter((r) => Number(draft[r.slug]) !== r.weight); + + const save = async () => { + setBusy(true); + setNotice(null); + setError(null); + try { + const weights = Object.fromEntries(changed.map((r) => [r.slug, Number(draft[r.slug])])); + const res = await fetch(`${origin}/admin/authority`, { + method: 'PUT', + headers: mutHeaders({ 'Content-Type': 'application/json' }), + body: JSON.stringify({ weights }), + }); + const body = await readJson(res); + if (!res.ok) throw new Error(apiError(body, res)); + setNotice(`${changed.length} project weight(s) saved; effective within 30 seconds`); + load(); + } catch (err) { + setError(String(err)); + } finally { + setBusy(false); + } + }; + + const shown = rows + .filter((r) => !filter || r.slug.toLowerCase().includes(filter.toLowerCase())) + .sort((a, b) => Number(b.weight !== 1) - Number(a.weight !== 1) || a.slug.localeCompare(b.slug)); + + return ( +
+ {error ? : null} + {notice ? : null} +

+ Above 1 never changes the answer itself: the project's closest passage is added to the references + when it is reasonably relevant (1.15 is strict, 1.3 is looser). Below 1 pushes a project down in the + ranking and 0 hides it. Questions that name a project are not affected. +

+ setFilter(e.target.value)} + />{' '} + + + + + + + + + + {shown.map((r) => ( + + + + + ))} + +
projectweight
{r.slug} + setDraft({ ...draft, [r.slug]: e.target.value })} + /> +
+
+ ); +} + function ScheduleTab({ origin }: { origin: string }) { const [jobs, setJobs] = useState([]); const [draft, setDraft] = useState>({}); diff --git a/application/frontend/src/pages/chatbot/chatbot.scss b/application/frontend/src/pages/chatbot/chatbot.scss index 582f786e3..a2e3b7e11 100644 --- a/application/frontend/src/pages/chatbot/chatbot.scss +++ b/application/frontend/src/pages/chatbot/chatbot.scss @@ -70,11 +70,9 @@ h1.ui.header.chatbot-title { margin-bottom: 1rem !important; } -.chatbot-toolbar { - display: flex; - justify-content: center; - margin: 0 auto 0.5rem; - max-width: 960px; +.chat-continue { + margin: 0.25rem 0 0.5rem; + text-align: left; } /* Tablet & below */ @@ -477,3 +475,26 @@ h1.ui.header.chatbot-title { transform: translateX(-50%) translateY(0); } } + +.rate-answer { + display: flex; + gap: 0.25rem; + margin-top: 0.5rem; + + .rate-btn { + background: none; + border: 1px solid transparent; + border-radius: 4px; + cursor: pointer; + padding: 0.2rem 0.4rem; + + &:hover { + border-color: #c8cdd3; + } + + &.is-selected { + background: #e8f0fe; + border-color: #7ba3f0; + } + } +} diff --git a/application/frontend/src/pages/chatbot/chatbot.test.tsx b/application/frontend/src/pages/chatbot/chatbot.test.tsx new file mode 100644 index 000000000..8438c2221 --- /dev/null +++ b/application/frontend/src/pages/chatbot/chatbot.test.tsx @@ -0,0 +1,166 @@ +import { fireEvent, render, waitFor } from '@testing-library/react'; +import React from 'react'; +import { MemoryRouter } from 'react-router-dom'; + +import { Chatbot, linkifyCitations } from './chatbot'; + +// The syntax highlighter ships ESM only; the thumbs tests do not render code blocks. +jest.mock('react-syntax-highlighter', () => ({ Prism: () => null })); +jest.mock('react-syntax-highlighter/dist/esm/styles/prism', () => ({ oneLight: {} })); +jest.mock('../../hooks/useEnvironment', () => ({ + useEnvironment: () => ({ name: 'test', apiUrl: '/rest/v1' }), +})); + +function res(body: unknown, status = 200) { + return Promise.resolve({ + ok: status >= 200 && status < 300, + status, + statusText: 'OK', + text: () => Promise.resolve(typeof body === 'string' ? body : JSON.stringify(body)), + json: () => Promise.resolve(body), + }); +} + +function setup(answerId: string | undefined, feedbackStatus = 200) { + const calls: { url: string; init?: any }[] = []; + (global as any).fetch = jest.fn((url: string, init?: any) => { + calls.push({ url: String(url), init }); + if (String(url).endsWith('/auth/user')) return res('someone'); + if (String(url).endsWith('/completion')) { + return res({ + response: 'SAMM has 3 maturity levels per practice.', + accurate: true, + answer_id: answerId, + }); + } + if (String(url).endsWith('/answer_feedback')) return res({ rating: 'up' }, feedbackStatus); + return res({}); + }); + const utils = render( + + + + ); + const ask = async () => { + fireEvent.change(utils.getByPlaceholderText(/Type your infosec question/), { + target: { value: 'levels in SAMM?' }, + }); + fireEvent.click(utils.getByText('Ask')); + await utils.findByText(/3 maturity levels/); + }; + return { ...utils, calls, ask }; +} + +describe('Chatbot thumbs up/down', () => { + afterEach(() => jest.clearAllMocks()); + + it('sends the rating for the answer and highlights the choice', async () => { + const { ask, findByLabelText, calls } = setup('ans-1'); + await ask(); + const up = await findByLabelText('Good answer'); + fireEvent.click(up); + await waitFor(() => expect(up.getAttribute('aria-pressed')).toBe('true')); + const post = calls.find((c) => c.url.endsWith('/answer_feedback')); + expect(JSON.parse(post!.init.body)).toEqual({ answer_id: 'ans-1', rating: 'up' }); + }); + + it('undoes the highlight when the server rejects the rating', async () => { + const { ask, findByLabelText } = setup('ans-2', 500); + await ask(); + const down = await findByLabelText('Bad answer'); + fireEvent.click(down); + await waitFor(() => expect(down.getAttribute('aria-pressed')).toBe('false')); + }); + + it('shows no thumbs and keeps the privacy paragraph when the answer was not logged', async () => { + const { ask, queryByLabelText, queryByText } = setup(undefined); + await ask(); + expect(queryByLabelText('Good answer')).toBeNull(); + expect(queryByText(/Privacy & Security/)).not.toBeNull(); + }); + + it('hides the "not stored" privacy paragraph once answers are stored', async () => { + const { ask, findByText, queryByText } = setup('ans-3'); + expect(queryByText(/Privacy & Security/)).not.toBeNull(); + await ask(); + await findByText(/3 maturity levels/); + expect(queryByText(/Privacy & Security/)).toBeNull(); + }); +}); + +describe('linkifyCitations', () => { + const urls = ['https://github.com/OWASP/a.md', 'https://github.com/OWASP/b.md']; + + it('turns [n] markers into links to the n-th source', () => { + expect(linkifyCitations('Imperfect [1][2].', urls)).toBe( + 'Imperfect [[1]](https://github.com/OWASP/a.md)[[2]](https://github.com/OWASP/b.md).' + ); + }); + + it('leaves markers without a matching source, and unsafe links, alone', () => { + expect(linkifyCitations('See [3] and [1].', ['https://x.org/a'])).toBe( + 'See [3] and [[1]](https://x.org/a).' + ); + expect(linkifyCitations('See [1].', ['javascript:alert(1)'])).toBe('See [1].'); + }); + + it('returns the text unchanged when there are no sources', () => { + expect(linkifyCitations('Plain [1] text', undefined)).toBe('Plain [1] text'); + expect(linkifyCitations('Plain [1] text', [])).toBe('Plain [1] text'); + }); +}); + +describe('Chatbot CRE reference', () => { + afterEach(() => jest.clearAllMocks()); + + it('links a CRE reference to its /cre/ page', async () => { + (global as any).fetch = jest.fn((url: string) => { + if (String(url).endsWith('/auth/user')) return res('someone'); + if (String(url).endsWith('/completion')) { + return res({ + response: 'Escape output.', + accurate: true, + table: [{ doctype: 'CRE', id: '366-835', name: 'Escape output against XSS' }], + }); + } + return res({}); + }); + const utils = render( + + + + ); + fireEvent.change(utils.getByPlaceholderText(/Type your infosec question/), { + target: { value: 'how do I prevent xss' }, + }); + fireEvent.click(utils.getByText('Ask')); + const link = (await utils.findByText(/366-835/)).closest('a')!; + expect(link.getAttribute('href')).toBe('/cre/366-835'); + expect(document.body.innerHTML).not.toContain('sectionid/undefined'); + }); +}); + +describe('Chatbot continue conversation', () => { + afterEach(() => jest.clearAllMocks()); + + it('sends only the new question unless "Continue conversation" is ticked', async () => { + const { ask, calls, container, getByText, getByPlaceholderText, findAllByText } = setup('ans-1'); + await ask(); + const send = async (text: string, expected: number) => { + fireEvent.change(getByPlaceholderText(/Type your infosec question/), { target: { value: text } }); + fireEvent.click(getByText('Ask')); + await findAllByText(/3 maturity levels/).then((els) => expect(els.length).toBeGreaterThan(0)); + await waitFor(() => expect(calls.filter((c) => c.url.endsWith('/completion'))).toHaveLength(expected)); + }; + await send('second question', 2); + const bodies = () => + calls.filter((c) => c.url.endsWith('/completion')).map((c) => JSON.parse(c.init.body)); + expect(bodies()[1].messages).toHaveLength(1); + fireEvent.click(container.querySelector('input[type=checkbox]')!); + await send('third question', 3); + expect(bodies()[2].messages).toHaveLength(3); // last question + its answer + the new question + await send('fourth question', 4); + expect(bodies()[3].messages.map((m: any) => m.content)).not.toContain('levels in SAMM?'); + expect(bodies()[3].messages).toHaveLength(3); + }); +}); diff --git a/application/frontend/src/pages/chatbot/chatbot.tsx b/application/frontend/src/pages/chatbot/chatbot.tsx index 4d001469e..eba9e6238 100644 --- a/application/frontend/src/pages/chatbot/chatbot.tsx +++ b/application/frontend/src/pages/chatbot/chatbot.tsx @@ -5,7 +5,7 @@ import { marked } from 'marked'; import React, { useEffect, useRef, useState } from 'react'; import { Prism as SyntaxHighlighter } from 'react-syntax-highlighter'; import { oneLight } from 'react-syntax-highlighter/dist/esm/styles/prism'; -import { Button, Container, Form, GridRow, Header, Icon } from 'semantic-ui-react'; +import { Button, Checkbox, Container, Form, GridRow, Header, Icon } from 'semantic-ui-react'; import { Grid } from 'semantic-ui-react'; import { useEnvironment } from '../../hooks'; @@ -17,6 +17,23 @@ export function renderChatMarkdown(markdown: string): string { return DOMPurify.sanitize(String(parsed), { USE_PROFILES: { html: true } }); } +/** Turn [n] markers into links to the n-th cited source (the backend numbers them to match the list). */ +export function linkifyCitations(markdown: string, urls: string[] | undefined): string { + if (!urls || urls.length === 0) return markdown; + const safe = (url: string) => { + try { + const protocol = new URL(url).protocol; + return protocol === 'http:' || protocol === 'https:'; + } catch { + return false; + } + }; + return markdown.replace(/\[(\d+)\]/g, (marker, num) => { + const url = urls[Number(num) - 1]; + return url && safe(url) ? `[${marker}](${url})` : marker; + }); +} + export const Chatbot = () => { type chatMessage = { timestamp: string; @@ -25,6 +42,8 @@ export const Chatbot = () => { data: Document[] | null; accurate: boolean; owaspCitations?: string[]; + answerId?: string; + rating?: 'up' | 'down'; }; interface ChatState { @@ -40,6 +59,7 @@ export const Chatbot = () => { const [error, setError] = useState(''); const [chat, setChat] = useState(DEFAULT_CHAT_STATE); const [user, setUser] = useState(''); + const [continueConversation, setContinueConversation] = useState(false); const [modelName, setModelName] = useState(''); function getModelDisplayName(modelName: string): string { @@ -56,6 +76,8 @@ export const Chatbot = () => { } const hasMessages = chatMessages.length > 0; + // Thumbs only appear when the deployment stores answers; the "not stored" privacy paragraph is then hidden. + const feedbackOn = chatMessages.some((m) => m.answerId); const messagesEndRef = useRef(null); const messagesContainerRef = useRef(null); const [showScrollToBottom, setShowScrollToBottom] = useState(false); @@ -110,7 +132,7 @@ export const Chatbot = () => { } }, []); - function processResponse(response: string) { + function processResponse(response: string, citations?: string[]) { const responses = response.split('```'); const res: React.ReactElement[] = []; @@ -120,7 +142,7 @@ export const Chatbot = () => {

); @@ -136,14 +158,29 @@ export const Chatbot = () => { return res; } - const HISTORY_CAP = 20; - - function startNewChat() { - setChatMessages([]); - setError(''); - setChat(DEFAULT_CHAT_STATE); - setModelName(''); - setShowScrollToBottom(false); + // "Continue conversation" sends only the last question and its answer. + const HISTORY_MESSAGES = 2; + + /** Thumbs up/down: show the choice at once, tell the server, and undo it if that fails. */ + function rateAnswer(idx: number, rating: 'up' | 'down') { + const answerId = chatMessages[idx]?.answerId; + if (!answerId) return; + const previous = chatMessages[idx].rating; + const setRating = (value?: 'up' | 'down') => + setChatMessages((prev) => prev.map((m, i) => (i === idx ? { ...m, rating: value } : m))); + setRating(rating); + fetch(`${apiUrl}/answer_feedback`, { + method: 'POST', + headers: { 'Content-Type': 'application/json', Accept: 'application/json' }, + body: JSON.stringify({ answer_id: answerId, rating }), + }) + .then((res) => { + if (!res.ok) throw new Error(`rating failed (${res.status})`); + }) + .catch((err) => { + console.error(err); + setRating(previous); + }); } async function onSubmit() { @@ -154,7 +191,7 @@ export const Chatbot = () => { setChat({ ...chat, term: '' }); setLoading(true); - const prior = chatMessages.slice(-(HISTORY_CAP - 1)).map((m) => ({ + const prior = (continueConversation ? chatMessages.slice(-HISTORY_MESSAGES) : []).map((m) => ({ role: m.role === 'assistant' ? 'assistant' : 'user', content: m.message, })); @@ -215,6 +252,7 @@ export const Chatbot = () => { : data.table, accurate: data.accurate, owaspCitations: Array.isArray(data?.owasp_agent?.citations) ? data.owasp_agent.citations : [], + answerId: typeof data?.answer_id === 'string' ? data.answer_id : undefined, }, ]); }) @@ -233,6 +271,22 @@ export const Chatbot = () => { return null; } + if (dt === 'cre' && d.id) { + // A Common Requirement: link to its own page (the /node/... form is for standard sections). + return ( +

+ ); + } + let link = `/node/${d.doctype.toLowerCase()}/${d.name}`; link += d.section ? `/section/${d.section}` : `/sectionid/${d.sectionID}`; @@ -256,25 +310,29 @@ export const Chatbot = () => { } function displayOwaspCitations(urls: string[]) { - const safe = urls.filter((url) => { + const isSafe = (url: string) => { try { const parsed = new URL(url); return parsed.protocol === 'http:' || parsed.protocol === 'https:'; } catch { return false; } - }); - if (!safe.length) return null; + }; + if (!urls.some(isSafe)) return null; return (
OWASP community sources
- {safe.map((url, i) => ( - - ))} + {urls.map((url, i) => + isSafe(url) ? ( + // The number is the position in the full list, so [n] in the answer always matches. +
+ [{i + 1}] + + {url} + +
+ ) : null + )}
); } @@ -288,13 +346,6 @@ export const Chatbot = () => {
OWASP OpenCRE Chat
- {hasMessages ? ( -
- -
- ) : null}
@@ -315,7 +366,30 @@ export const Chatbot = () => { {m.timestamp}
-
{processResponse(m.message)}
+
{processResponse(m.message, m.owaspCitations)}
+ + {m.role === 'assistant' && m.answerId ? ( +
+ + +
+ ) : null} {m.data && m.data.length > 0 && (
@@ -373,6 +447,14 @@ export const Chatbot = () => { onChange={(e) => setChat({ ...chat, term: e.target.value })} placeholder="Type your infosec question here…" /> + {hasMessages ? ( + setContinueConversation(!!d.checked)} + /> + ) : null} @@ -392,11 +474,15 @@ export const Chatbot = () => { Model operation is generously sponsored by{' '} Software Improvement Group.
-
- Privacy & Security: Your question is sent to Heroku, the hosting provider for OpenCRE, and - then to GCP, all via protected connections. Your data isn't stored on OpenCRE servers. The - OpenCRE team employed extensive measures to ensure privacy and security. To review the code: - https://github.com/owasp/OpenCRE + {!feedbackOn && ( + <> +
+ Privacy & Security: Your question is sent to Heroku, the hosting provider for OpenCRE, and + then to GCP, all via protected connections. Your data isn't stored on OpenCRE servers. The + OpenCRE team employed extensive measures to ensure privacy and security. To review the + code: https://github.com/owasp/OpenCRE + + )}
diff --git a/application/frontend/src/utils/document.ts b/application/frontend/src/utils/document.ts index bf01714b5..800f699c5 100644 --- a/application/frontend/src/utils/document.ts +++ b/application/frontend/src/utils/document.ts @@ -32,7 +32,7 @@ export const groupLinksByType = (node: Document): LinksByType => node.links ? groupBy(node.links, (link) => link.ltype) : {}; export const orderLinksByType = (lbt: LinksByType): LinksByType => { - const order = ['Linked To', 'Automatically linked to', 'Contains', 'Is Part Of', 'Related']; + const order = ['Linked To', 'Automatically linked to', 'Contains', 'Related', 'Is Part Of']; const res: LinksByType = {}; for (const itm of order) { if (lbt[itm]) { diff --git a/application/mcp/ask.py b/application/mcp/ask.py index 7030368ab..be24a13b4 100644 --- a/application/mcp/ask.py +++ b/application/mcp/ask.py @@ -10,6 +10,7 @@ logger = get_logger(__name__) +import time from typing import Any, Dict, List, Mapping, Optional, Set, Tuple from application.mcp.rest_client import ( @@ -61,6 +62,21 @@ "type": "string", "description": "Optional place hint for geo-sensitive meta questions", }, + "history": { + "type": "array", + "description": ( + "Optional earlier turns [{role: user|assistant, content}] so follow-up " + "questions keep their context (the web chat sends the same)." + ), + "items": { + "type": "object", + "properties": { + "role": {"type": "string", "enum": ["user", "assistant"]}, + "content": {"type": "string"}, + }, + "required": ["role", "content"], + }, + }, }, "required": ["prompt"], "additionalProperties": False, @@ -68,6 +84,24 @@ } +RATE_TOOL: Dict[str, Any] = { + "name": "rate_answer", + "summary": ( + "Give a thumbs up or down to an earlier ask answer, using its answer_id. " + "Optional comment says what was wrong or good. Ratings are used to improve answers." + ), + "input_schema": { + "type": "object", + "properties": { + "answer_id": {"type": "string", "description": "answer_id returned by ask"}, + "rating": {"type": "string", "enum": ["up", "down"]}, + "comment": {"type": "string", "description": "Optional short reason"}, + }, + "required": ["answer_id", "rating"], + "additionalProperties": False, + }, +} + _APP = None @@ -82,28 +116,37 @@ def _ensure_app(): return _APP -def _router_handle(prompt: str) -> Optional[Dict[str, Any]]: - """Run OwaspAgentRouter inside an app context so the sim gate can use pgvector. +def _router_handle( + prompt: str, history: Optional[List[Dict[str, str]]] = None +) -> Optional[Dict[str, Any]]: + """Run the shared OWASP agent entry point inside an app context. - MCP stdio has no request context; unit tests may call without Flask — fall - back to a bare handle() if app bootstrap is unavailable. + MCP stdio has no request context; the sim gate needs one for pgvector. Raises + ``OwaspAgentUnavailable`` (same rule as the web chat) when the agent crashes on a meta question. """ + from application.utils.owasp_agent.service import route_question + try: from flask import has_app_context if has_app_context(): - return OwaspAgentRouter().handle(prompt) - except Exception: + return route_question(prompt, history) + except ImportError: pass - try: - app = _ensure_app() - with app.app_context(): - return OwaspAgentRouter().handle(prompt) - except Exception: - try: - return OwaspAgentRouter().handle(prompt) - except Exception: - return None + app = _ensure_app() + with app.app_context(): + return route_question(prompt, history) + + +def _in_app_context(fn: Any) -> Any: + """Run ``fn`` inside a Flask app context (MCP has none); None if the app cannot start.""" + from flask import has_app_context + + if has_app_context(): + return fn() + app = _ensure_app() + with app.app_context(): + return fn() def ask( @@ -111,6 +154,42 @@ def ask( *, corpus: str = CORPUS_AUTO, place: Optional[str] = None, + history: Optional[List[Dict[str, str]]] = None, + rest_client: Optional[RestClient] = None, +) -> Dict[str, Any]: + """Route ``prompt`` and return a fact pack; the pack carries an ``answer_id`` for rating.""" + started = time.monotonic() + result = _ask_impl( + prompt, corpus=corpus, place=place, history=history, rest_client=rest_client + ) + try: + from application.utils.answer_feedback import attach_answer_id + + elapsed = int((time.monotonic() - started) * 1000) + return _in_app_context( + lambda: attach_answer_id( + result, question=prompt, source="mcp", latency_ms=elapsed + ) + ) + except Exception: # noqa: BLE001 -- rating support must never break answering + return result + + +def rate_answer( + answer_id: str, rating: str, comment: Optional[str] = None +) -> Dict[str, Any]: + """Record a thumbs up/down for an earlier answer.""" + from application.utils.answer_feedback import record_rating + + return _in_app_context(lambda: record_rating(answer_id, rating, comment)) + + +def _ask_impl( + prompt: str, + *, + corpus: str = CORPUS_AUTO, + place: Optional[str] = None, + history: Optional[List[Dict[str, str]]] = None, rest_client: Optional[RestClient] = None, ) -> Dict[str, Any]: """Route ``prompt`` and return a fact pack for the chosen corpus.""" @@ -127,16 +206,23 @@ def ask( if place_text and corpus_key in (CORPUS_AUTO, CORPUS_OWASP_META): effective = f"{text} near {place_text}" + from application.utils.chat_history import normalize_chat_messages + + history = normalize_chat_messages(history) or None intent = classify_intent(effective) if corpus_key == CORPUS_SECURITY: return _security_facts(text, rest_client=rest_client, intent=intent) if corpus_key == CORPUS_OWASP_META: - return _owasp_meta_facts(effective, intent=intent, allow_cre_defer=False) + return _owasp_meta_facts( + effective, intent=intent, allow_cre_defer=False, history=history + ) # auto: meta answer wins when the router handled it; otherwise CRE facts. - meta = _owasp_meta_facts(effective, intent=intent, allow_cre_defer=True) + meta = _owasp_meta_facts( + effective, intent=intent, allow_cre_defer=True, history=history + ) if meta.get("defer_to_security") or meta.get("channel") == "none": return _security_facts(text, rest_client=rest_client, intent=intent) return meta @@ -147,6 +233,7 @@ def _owasp_meta_facts( *, intent: str, allow_cre_defer: bool, + history: Optional[List[Dict[str, str]]] = None, ) -> Dict[str, Any]: if not is_owasp_agent_enabled(): if allow_cre_defer: @@ -169,7 +256,18 @@ def _owasp_meta_facts( # Always try the meta/sim router first — evidence and community cards can # answer prompts that classify_intent labels as cre_normative. - resp = _router_handle(prompt) + from application.utils.owasp_agent.service import OwaspAgentUnavailable + + try: + resp = _router_handle(prompt, history) + except OwaspAgentUnavailable: + # Same rule as the web chat: never answer a community/meta question from CRE retrieval. + return { + "ok": False, + "channel": CORPUS_OWASP_META, + "intent": intent, + "message": "The OWASP agent hit a temporary problem. Please try again.", + } if resp is None: if intent == "cre_normative" and not allow_cre_defer: return { diff --git a/application/mcp/remote_auth.py b/application/mcp/remote_auth.py index 52de9dab5..f3b61cf18 100644 --- a/application/mcp/remote_auth.py +++ b/application/mcp/remote_auth.py @@ -49,37 +49,70 @@ def extract_bearer(authorization: Optional[str] = None) -> Optional[str]: return None -def require_remote_ask_auth() -> ApiKeyIdentity: - """Verify Bearer Admin JWT + allowlist. Raises RemoteAskAuthError on failure.""" +# Minimum role per MCP tool. Everything we ship today is read-only or answers questions, so any +# account with a role may call it. Tools that change state or expose operations go in +# ADMIN_ONLY_TOOLS (named admins only). Unknown tool names default to admin: deny by default. +ADMIN_ONLY_TOOLS: frozenset = frozenset() +USER_TOOL_PREFIXES = ("owasp_meta_",) +USER_TOOLS = frozenset({"ask", "rate_answer"}) + + +def min_role_for_tool(tool_name: str, public_cre_tools=frozenset()) -> Optional[str]: + """Role a tool needs; ``None`` for the public, read-only CRE REST tools (as before).""" + from application.utils import roles + + if tool_name in ADMIN_ONLY_TOOLS: + return roles.ADMIN + if tool_name in USER_TOOLS or tool_name.startswith(USER_TOOL_PREFIXES): + return roles.USER + if tool_name in public_cre_tools: + return None + return roles.ADMIN + + +def require_remote_auth( + tool_name: str, public_cre_tools=frozenset() +) -> Optional[ApiKeyIdentity]: + """Verify the Bearer API key and that its owner holds the role ``tool_name`` needs. + + Returns ``None`` for the public CRE read tools (no key needed). + """ + from application.utils import roles + from application.utils.api_tokens.verify import ( + verify_api_token, + verify_token_with_role, + ) + + needed = min_role_for_tool(tool_name, public_cre_tools) + if needed is None: + return None bearer = extract_bearer() if not bearer: raise RemoteAskAuthError( - 401, - "Unauthorized: Bearer Admin API token required for remote ask.", + 401, "Unauthorized: Bearer API key required (sign in and create one)." ) - identity = verify_admin_allowlisted_api_token(bearer) + identity = verify_token_with_role(bearer, needed) if identity is None: - # Distinguish invalid/revoked vs allowlist via a second cheap check. - from application.utils.api_tokens.verify import ( - email_on_login_allowlist, - verify_api_token, - ) - base = verify_api_token(bearer) if base is None: raise RemoteAskAuthError( - 401, - "Unauthorized: invalid, expired, or revoked API token.", + 401, "Unauthorized: invalid, expired, or revoked API key." ) - if not email_on_login_allowlist(base.email): + if roles.role_for(base.email) is None: raise RemoteAskAuthError( - 403, - "Forbidden: API key owner is not on the allowlist for this environment.", + 403, "Forbidden: this account has no access to this environment." ) - raise RemoteAskAuthError(401, "Unauthorized: invalid API token.") + raise RemoteAskAuthError( + 403, f"Forbidden: the {needed} role is required for the tool '{tool_name}'." + ) try: bind_api_key_identity(identity) except Exception: # Binding last_used_at is best-effort; auth already succeeded. pass return identity + + +def require_remote_ask_auth() -> ApiKeyIdentity: + """Back-compat name: the ``ask`` tool needs the user role.""" + return require_remote_auth("ask") diff --git a/application/mcp/server.py b/application/mcp/server.py index 50d5ad373..2039a13bc 100644 --- a/application/mcp/server.py +++ b/application/mcp/server.py @@ -14,7 +14,7 @@ from mcp.server.lowlevel import Server from mcp.server.stdio import stdio_server -from application.mcp.ask import ASK_TOOL, ask as run_ask +from application.mcp.ask import ASK_TOOL, RATE_TOOL, ask as run_ask, rate_answer from application.mcp.catalog import active_tools, auth_tools_enabled from application.mcp.openapi_loader import input_schemas_for from application.mcp.owasp_agent_tools import ( @@ -35,6 +35,7 @@ SERVER_NAME = "opencre" SERVER_VERSION = "0.1.0" ASK_TOOL_NAME = ASK_TOOL["name"] +RATE_TOOL_NAME = RATE_TOOL["name"] def server_instructions() -> str: @@ -170,22 +171,34 @@ async def on_call_tool( name = params.name arguments: Dict[str, Any] = dict(params.arguments or {}) try: - if name == ASK_TOOL_NAME: - from application.mcp import remote_auth + from application.mcp import remote_auth + + if remote_auth.remote_mode_enabled(): + # Every tool needs a signed-in account with the role it requires. + try: + caller = remote_auth.require_remote_auth(name, frozenset(schemas)) + except remote_auth.RemoteAskAuthError as auth_exc: + return types.CallToolResult( + content=[types.TextContent(type="text", text=str(auth_exc))], + is_error=True, + ) + if name == ASK_TOOL_NAME and caller is not None: + from application.utils import rate_limit, roles - if remote_auth.remote_mode_enabled(): try: - remote_auth.require_remote_ask_auth() - except remote_auth.RemoteAskAuthError as auth_exc: + rate_limit.check_and_count( + caller.email, roles.role_for(caller.email), "ask" + ) + except rate_limit.RateLimited as limited: return types.CallToolResult( content=[ types.TextContent( - type="text", - text=str(auth_exc), + type="text", text=f"429 {limited.message()}" ) ], is_error=True, ) + if name == ASK_TOOL_NAME: data = await run_blocking_with_app( lambda: run_ask( str(arguments.get("prompt") or ""), @@ -195,6 +208,7 @@ async def on_call_tool( if arguments.get("place") is not None else None ), + history=arguments.get("history"), rest_client=client, ) ) @@ -206,6 +220,17 @@ async def on_call_tool( ), is_error=False, ) + if name == RATE_TOOL_NAME: + data = rate_answer( + str(arguments.get("answer_id") or ""), + str(arguments.get("rating") or ""), + str(arguments["comment"]) if arguments.get("comment") else None, + ) + return types.CallToolResult( + content=[types.TextContent(type="text", text=json.dumps(data))], + structured_content=data, + is_error=False, + ) if get_owasp_meta_tool(name) is not None: if name.startswith("owasp_meta_sources_"): data = await run_blocking_with_app( @@ -307,6 +332,12 @@ def dispatch_tool( place=str(args["place"]) if args.get("place") is not None else None, rest_client=rest_client, ) + if tool_name == RATE_TOOL_NAME: + return rate_answer( + str(args.get("answer_id") or ""), + str(args.get("rating") or ""), + str(args["comment"]) if args.get("comment") else None, + ) if tool_name in list_owasp_meta_tool_names() or get_owasp_meta_tool(tool_name): return call_owasp_meta_tool(tool_name, args) client = rest_client or RestClient() diff --git a/application/prompt_client/litellm_router.py b/application/prompt_client/litellm_router.py index 53f3c6614..78d19b6df 100644 --- a/application/prompt_client/litellm_router.py +++ b/application/prompt_client/litellm_router.py @@ -188,9 +188,31 @@ def system_user_fn( *, temperature: float = 0.0, extra_try_kwargs: Optional[dict[str, Any]] = None, + interactive: bool = False, + fallback_model: Optional[str] = None, **kwargs: Any, ) -> LlmFn: - """Build ``(system, user) -> text`` for Librarian / OIE call sites.""" + """Build ``(system, user) -> text`` for Librarian / OIE call sites. + + ``interactive=True`` is for calls a person is waiting on: quick retries, ``fallback_model`` and the + circuit breaker (see :func:`resilient_completion`) instead of the long bulk-job sleeps. + """ + + def _text(messages: list, **extra: Any) -> str: + if interactive: + return extract_content_text( + resilient_completion( + model=model, + messages=messages, + fallback_model=fallback_model, + temperature=temperature, + **extra, + **kwargs, + ) + ) + return completion_text( + model=model, messages=messages, temperature=temperature, **extra, **kwargs + ) def _call(system: str, user: str) -> str: messages = [ @@ -199,24 +221,134 @@ def _call(system: str, user: str) -> str: ] if extra_try_kwargs: try: - return completion_text( - model=model, - messages=messages, - temperature=temperature, - **extra_try_kwargs, - **kwargs, - ) - except Exception: + return _text(messages, **extra_try_kwargs) + except Exception as err: # noqa: BLE001 + if interactive and is_transient_llm_error(err): + raise # an outage is not a bad-kwargs problem; do not repeat the whole retry ladder logger.debug( "LiteLLM extra kwargs failed for model=%s; retrying without them", model, exc_info=True, ) - return completion_text( - model=model, - messages=messages, - temperature=temperature, - **kwargs, - ) + return _text(messages) return _call + + +# --- Interactive calls (chat / MCP ask): quick retries, a fallback model and a circuit breaker ------------- +# +# Bulk jobs (ingestion, Module B/C) keep ``with_rate_limit_retry``: long fixed sleeps are fine there. A person +# waiting for an answer needs short, jittered back-off, a cheaper model to fall back on during a provider +# outage, and fast failure once the primary model is clearly down. + +import random +import threading + +_BREAKER_LOCK = threading.Lock() +_BREAKER: dict = {} # model -> {"fails": int, "open_until": float} + + +def interactive_policy() -> tuple[int, float, float]: + """(retries, first delay seconds, max delay seconds) for interactive calls.""" + return ( + int(os.environ.get("LLM_INTERACTIVE_RETRIES", "3")), + float(os.environ.get("LLM_INTERACTIVE_BASE_DELAY", "1.0")), + float(os.environ.get("LLM_INTERACTIVE_MAX_DELAY", "8.0")), + ) + + +def _breaker_open(model: str) -> bool: + with _BREAKER_LOCK: + state = _BREAKER.get(model) + return bool(state and state["open_until"] > time.time()) + + +def _breaker_record(model: str, ok: bool) -> None: + threshold = int(os.environ.get("LLM_BREAKER_FAILURES", "3")) + cooldown = float(os.environ.get("LLM_BREAKER_COOLDOWN_SECONDS", "30")) + with _BREAKER_LOCK: + state = _BREAKER.setdefault(model, {"fails": 0, "open_until": 0.0}) + if ok: + state["fails"], state["open_until"] = 0, 0.0 + return + state["fails"] += 1 + if state["fails"] >= threshold: + state["open_until"] = time.time() + cooldown + state["fails"] = 0 + logger.warning("LLM circuit open for %s for %.0fs", model, cooldown) + + +def reset_breaker() -> None: + with _BREAKER_LOCK: + _BREAKER.clear() + + +def quick_retry(fn: Callable[[], Any], *, context: str) -> Any: + """Retry transient errors with exponential back-off and jitter (interactive profile).""" + retries, base, cap = interactive_policy() + for attempt in range(retries + 1): + try: + return fn() + except Exception as err: # noqa: BLE001 + if attempt >= retries or not is_transient_llm_error(err): + raise + delay = min(cap, base * (2**attempt)) * (0.75 + random.random() * 0.5) + logger.info( + "%s: transient %s, retrying in %.1fs (%s/%s)", + context, + type(err).__name__, + delay, + attempt + 1, + retries, + ) + time.sleep(delay) + raise RuntimeError("unreachable: retry loop exited unexpectedly") + + +def resilient_completion( + *, + model: str, + messages: Sequence[dict[str, Any]], + fallback_model: Optional[str] = None, + context: str = "LiteLLM completion", + client: Any = None, + **kwargs: Any, +) -> Any: + """Completion for a waiting user: quick retries on ``model``, then ``fallback_model``. + + A model that keeps failing trips a short circuit breaker so the next requests go straight to the + fallback instead of each waiting through the retries. Non-transient errors (bad key, refusal, + bad request) are raised immediately and never trigger a fallback. + """ + llm = client if client is not None else get_litellm() + kwargs.setdefault( + "timeout", float(os.environ.get("LLM_INTERACTIVE_TIMEOUT_SECONDS", "90")) + ) + + def _on(m: str) -> Any: + return quick_retry( + lambda: llm.completion(model=m, messages=list(messages), **kwargs), + context=f"{context} [{m}]", + ) + + if not _breaker_open(model): + try: + out = _on(model) + _breaker_record(model, True) + return out + except Exception as err: # noqa: BLE001 + if not is_transient_llm_error(err): + raise + _breaker_record(model, False) + if not fallback_model or fallback_model == model: + raise + logger.warning( + "%s: %s failed (%s); falling back to %s", + context, + model, + type(err).__name__, + fallback_model, + ) + elif not fallback_model or fallback_model == model: + raise RuntimeError(f"model {model} is temporarily unavailable (circuit open)") + return _on(fallback_model) # type: ignore[arg-type] diff --git a/application/prompt_client/llm_error_utils.py b/application/prompt_client/llm_error_utils.py index 32dc724a9..d5fbed0c1 100644 --- a/application/prompt_client/llm_error_utils.py +++ b/application/prompt_client/llm_error_utils.py @@ -33,9 +33,51 @@ def is_rate_limit_error(err: BaseException) -> bool: return False +_OVERLOAD_NAMES = ( + "serviceunavailable", + "internalservererror", + "badgateway", + "gatewaytimeout", + "overloaded", +) +_OVERLOAD_MESSAGES = ( + "service unavailable", + "overloaded", + "high demand", + "currently experiencing", + "try again later", + "internal error", + "internal server error", + "bad gateway", + "gateway timeout", + "model is unavailable", +) + + +def is_overload_error(err: BaseException) -> bool: + """Provider-side outage / overload (HTTP 500, 502, 503, 504): retry, then fall back.""" + name = type(err).__name__.lower() + if any(tok in name for tok in _OVERLOAD_NAMES): + return True + status = ( + getattr(err, "status", None) + or getattr(err, "status_code", None) + or getattr(err, "http_status", None) + ) + if status in (500, 502, 503, 504): + return True + msg = str(err).lower() + if any(tok in msg for tok in _OVERLOAD_MESSAGES): + return True + return any( + f'"code": {code}' in msg or f"error {code}" in msg + for code in (500, 502, 503, 504) + ) + + def is_transient_llm_error(err: BaseException) -> bool: - """Connection / timeout blips that are safe to retry (not model refusal).""" - if is_rate_limit_error(err): + """Connection / timeout blips and provider overloads that are safe to retry (not model refusal).""" + if is_rate_limit_error(err) or is_overload_error(err): return True name = type(err).__name__.lower() if any( diff --git a/application/tests/alembic_single_head_test.py b/application/tests/alembic_single_head_test.py new file mode 100644 index 000000000..8a482f14b --- /dev/null +++ b/application/tests/alembic_single_head_test.py @@ -0,0 +1,37 @@ +"""The migration chain must end in exactly one head. + +Two heads break `flask db upgrade` ("Multiple head revisions"). When two branches each add +a migration, add a merge migration whose down_revision lists both parents. +""" + +import re +import unittest +from pathlib import Path + +VERSIONS = Path(__file__).resolve().parents[2] / "migrations" / "versions" + + +def _heads() -> list: + revisions, parents = {}, set() + for path in VERSIONS.glob("*.py"): + text = path.read_text() + rev = re.search(r"^revision\s*=\s*['\"]([^'\"]+)", text, re.M) + down = re.search(r"^down_revision\s*=\s*(.+)$", text, re.M) + if rev: + revisions[rev.group(1)] = path.name + if down: + parents.update(re.findall(r"['\"]([0-9a-zA-Z_]+)['\"]", down.group(1))) + return sorted(r for r in revisions if r not in parents) + + +class AlembicSingleHeadTest(unittest.TestCase): + def test_exactly_one_head(self) -> None: + self.assertEqual( + len(_heads()), + 1, + f"migration chain has multiple heads: {_heads()}; add a merge migration", + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/application/tests/answer_feedback_test.py b/application/tests/answer_feedback_test.py new file mode 100644 index 000000000..49925cb00 --- /dev/null +++ b/application/tests/answer_feedback_test.py @@ -0,0 +1,183 @@ +import os +import unittest +from unittest import mock + +from application import create_app, sqla +from application.utils import answer_feedback as af + + +class AnswerFeedbackTest(unittest.TestCase): + def setUp(self) -> None: + patcher = mock.patch.dict(os.environ, {"NO_LOAD_GRAPH_DB": "1"}) + patcher.start() + self.addCleanup(patcher.stop) + self.app = create_app(mode="test") + self.ctx = self.app.app_context() + self.ctx.push() + sqla.create_all() + + def tearDown(self) -> None: + os.environ.pop("ANSWER_LOG_ENABLED", None) + os.environ.pop("ANSWER_LOG_CHAT", None) + sqla.session.remove() + sqla.drop_all() + self.ctx.pop() + + def test_chat_answers_are_not_logged_unless_the_deployment_opts_in(self) -> None: + self.assertIsNone(af.log_answer(source="chat", question="q", answer="a")) + self.assertTrue(af.log_answer(source="mcp", question="q", answer="a")) + os.environ["ANSWER_LOG_CHAT"] = "1" + self.assertTrue(af.log_answer(source="chat", question="q", answer="a")) + + def test_log_then_rate_and_change_rating(self) -> None: + aid = af.log_answer( + source="mcp", + question="How many levels in SAMM?", + answer="Three [1]", + citations=["https://x/a"], + channel="owasp_meta", + ) + self.assertTrue(aid) + self.assertEqual(af.record_rating(aid, "UP", " nice ")["rating"], "up") + self.assertEqual(af.record_rating(aid, "down")["rating"], "down") + rows = af.list_feedback() + self.assertEqual(len(rows), 1) + self.assertEqual( + (rows[0]["rating"], rows[0]["comment"], rows[0]["citations"]), + ("down", None, ["https://x/a"]), + ) + self.assertEqual(af.list_feedback(rating="up"), []) + + def test_rejects_bad_input(self) -> None: + aid = af.log_answer(source="mcp", question="q", answer="a") + with self.assertRaises(ValueError): + af.record_rating(aid, "meh") + with self.assertRaises(ValueError): + af.record_rating("", "up") + with self.assertRaises(LookupError): + af.record_rating("no-such-id", "up") + + def test_logging_can_be_switched_off_and_never_raises(self) -> None: + os.environ["ANSWER_LOG_ENABLED"] = "0" + self.assertIsNone(af.log_answer(source="mcp", question="q", answer="a")) + os.environ["ANSWER_LOG_ENABLED"] = "1" + sqla.drop_all() # table missing: must be swallowed, not raised + self.assertIsNone(af.log_answer(source="mcp", question="q", answer="a")) + sqla.create_all() + + def test_attach_answer_id_adds_id_only_when_logged(self) -> None: + out = af.attach_answer_id( + {"message": "m", "citations": ["c"], "channel": "x"}, + question="q", + source="mcp", + ) + self.assertIn("answer_id", out) + os.environ["ANSWER_LOG_ENABLED"] = "0" + self.assertNotIn( + "answer_id", + af.attach_answer_id({"message": "m"}, question="q", source="mcp"), + ) + + +if __name__ == "__main__": + unittest.main() + + +class McpRatingTest(unittest.TestCase): + def setUp(self) -> None: + patcher = mock.patch.dict(os.environ, {"NO_LOAD_GRAPH_DB": "1"}) + patcher.start() + self.addCleanup(patcher.stop) + self.app = create_app(mode="test") + self.ctx = self.app.app_context() + self.ctx.push() + sqla.create_all() + + def tearDown(self) -> None: + sqla.session.remove() + sqla.drop_all() + self.ctx.pop() + + def test_ask_result_carries_id_that_rate_answer_accepts(self) -> None: + from unittest import mock + + from application.mcp import ask as ask_mod + + fake = { + "ok": True, + "channel": "owasp_meta", + "message": "SAMM has 3 levels", + "citations": ["https://x"], + } + with mock.patch.object(ask_mod, "_ask_impl", return_value=fake): + out = ask_mod.ask("how many levels in samm?") + self.assertIn("answer_id", out) + self.assertEqual( + ask_mod.rate_answer(out["answer_id"], "up", "correct")["rating"], "up" + ) + with self.assertRaises(KeyError): + ask_mod.rate_answer("nope", "up") + + def test_rate_tool_is_listed_and_dispatches(self) -> None: + from application.mcp import server + from application.utils import answer_feedback as af_mod + + aid = af_mod.log_answer(source="mcp", question="q", answer="a") + out = server.dispatch_tool( + "rate_answer", {"answer_id": aid, "rating": "down", "comment": "wrong"} + ) + self.assertEqual(out["rating"], "down") + + +class RestFeedbackTest(unittest.TestCase): + def setUp(self) -> None: + patcher = mock.patch.dict(os.environ, {"NO_LOAD_GRAPH_DB": "1"}) + patcher.start() + self.addCleanup(patcher.stop) + self.app = create_app(mode="test") + self.ctx = self.app.app_context() + self.ctx.push() + sqla.create_all() + self.app.secret_key = "test-secret" + os.environ["INSECURE_REQUESTS"] = "True" + self.client = self.app.test_client() + + def tearDown(self) -> None: + sqla.session.remove() + sqla.drop_all() + self.ctx.pop() + + def _login(self) -> None: + os.environ["NO_LOGIN"] = "1" # same bypass the chat completion tests use + self.addCleanup(os.environ.pop, "NO_LOGIN", None) + + def test_with_answer_id_adds_id_for_dict_responses_only(self) -> None: + from application.web import web_main + + os.environ["ANSWER_LOG_CHAT"] = "1" + self.addCleanup(os.environ.pop, "ANSWER_LOG_CHAT", None) + out = web_main._with_answer_id( + {"response": "hi", "owasp_agent": {"citations": ["c"]}}, "q?" + ) + self.assertIn("answer_id", out) + self.assertEqual(web_main._with_answer_id("plain", "q"), "plain") + + def test_post_feedback_requires_valid_rating_and_known_id(self) -> None: + from application.utils import answer_feedback as af_mod + + aid = af_mod.log_answer(source="mcp", question="q", answer="a") + self._login() + bad = self.client.post( + "/rest/v1/answer_feedback", json={"answer_id": aid, "rating": "meh"} + ) + self.assertEqual(bad.status_code, 400) + missing = self.client.post( + "/rest/v1/answer_feedback", json={"answer_id": "nope", "rating": "up"} + ) + self.assertEqual(missing.status_code, 404) + ok = self.client.post( + "/rest/v1/answer_feedback", + json={"answer_id": aid, "rating": "up", "comment": "good"}, + ) + self.assertEqual(ok.status_code, 200) + self.assertEqual(ok.get_json()["rating"], "up") diff --git a/application/tests/chat_ask_parity_test.py b/application/tests/chat_ask_parity_test.py new file mode 100644 index 000000000..91a416d2c --- /dev/null +++ b/application/tests/chat_ask_parity_test.py @@ -0,0 +1,104 @@ +"""The web chatbot and MCP ask share one entry point: same routing, same answer, same failures.""" + +from __future__ import annotations + +import os +import unittest +from unittest.mock import MagicMock, patch + +from application import create_app, sqla +from application.mcp.ask import ask +from application.utils.owasp_agent.service import OwaspAgentUnavailable + +AGENT_RESPONSE = { + "response": "Answer: three maturity levels per practice [1].", + "accurate": True, + "owasp_agent": { + "ok": True, + "citations": ["https://github.com/OWASP/samm/blob/HEAD/a.md"], + "data": {"evidence": [{"n": 1, "path": "a.md", "snippet": "three levels"}]}, + }, +} + + +class ChatAskParityTest(unittest.TestCase): + def setUp(self) -> None: + self._env = patch.dict( + os.environ, + { + "NO_LOGIN": "1", + "NO_LOAD_GRAPH_DB": "1", + "OWASP_AGENT_ENABLED": "1", + "INSECURE_REQUESTS": "1", + }, + ) + self._env.start() + self.addCleanup(self._env.stop) + self.app = create_app(mode="test") + self.ctx = self.app.app_context() + self.ctx.push() + sqla.create_all() + + def tearDown(self) -> None: + sqla.session.remove() + sqla.drop_all() + self.ctx.pop() + + def _router(self, handle): + router = MagicMock() + router.handle.side_effect = handle + return ( + patch( + "application.utils.owasp_agent.OwaspAgentRouter", return_value=router + ), + router, + ) + + def test_same_question_same_answer_and_citations(self): + history = [ + {"role": "user", "content": "tell me about SAMM"}, + {"role": "assistant", "content": "ok"}, + ] + patcher, router = self._router(lambda prompt, history=None: AGENT_RESPONSE) + with patcher: + chat = ( + self.app.test_client() + .post( + "/rest/v1/completion", + json={ + "prompt": "How many levels does SAMM have?", + "messages": history, + }, + ) + .get_json() + ) + via_ask = ask("How many levels does SAMM have?", history=history) + self.assertEqual(chat["response"], via_ask["message"]) + self.assertEqual(chat["owasp_agent"]["citations"], via_ask["citations"]) + self.assertEqual( + via_ask["evidence"], AGENT_RESPONSE["owasp_agent"]["data"]["evidence"] + ) + # Both routed through the one router with the conversation history. + for call in router.handle.call_args_list: + self.assertEqual(call.kwargs.get("history"), history) + + def test_agent_failure_on_a_meta_question_is_refused_by_both(self): + def boom(prompt, history=None): + raise RuntimeError("router crashed") + + patcher, _ = self._router(boom) + with patcher: + chat = self.app.test_client().post( + "/rest/v1/completion", json={"prompt": "Who leads the Berlin chapter?"} + ) + via_ask = ask("Who leads the Berlin chapter?") + self.assertEqual(chat.status_code, 503) + self.assertFalse(via_ask["ok"]) + self.assertIn("temporary problem", via_ask["message"]) + + def test_service_returns_none_when_disabled(self): + from application.utils.owasp_agent.service import route_question + + with patch.dict(os.environ, {"OWASP_AGENT_ENABLED": "0"}): + self.assertIsNone(route_question("anything")) + self.assertTrue(issubclass(OwaspAgentUnavailable, Exception)) diff --git a/application/tests/chat_history_test.py b/application/tests/chat_history_test.py new file mode 100644 index 000000000..c5359272e --- /dev/null +++ b/application/tests/chat_history_test.py @@ -0,0 +1,171 @@ +import os +import unittest +from unittest import mock + +from application.utils.owasp_agent import router as r + +HISTORY = [ + {"role": "user", "content": "Does the exchange cover agentic AI?"}, + { + "role": "assistant", + "content": "Yes, the AI Exchange covers agentic AI threats and controls.", + }, +] + + +class NeedsContextTest(unittest.TestCase): + def test_follow_ups_need_context_full_questions_do_not(self): + for q in ( + "What are some key risks", + "and the controls?", + "Why is that?", + "tell me more", + "Does it cover prompt injection in production systems?", + ): + self.assertTrue(r.needs_context(q), q) + for q in ( + "What are the three maturity levels defined by OWASP SAMM for each practice?", + "Who leads the OWASP Developer Guide project and which companies sponsor it?", + ): + self.assertFalse(r.needs_context(q), q) + + +class StandaloneQuestionTest(unittest.TestCase): + def test_self_contained_question_ignores_history_and_never_calls_the_model(self): + llm = mock.Mock() + q = "What are the three maturity levels defined by OWASP SAMM for each practice?" + self.assertEqual(r.standalone_question(q, HISTORY, llm), q) + llm.assert_not_called() + + def test_follow_up_is_rewritten_and_does_not_contain_the_previous_question(self): + llm = mock.Mock( + return_value="What are some key risks of agentic AI according to the AI Exchange?" + ) + out = r.standalone_question("What are some key risks", HISTORY, llm) + self.assertEqual( + out, "What are some key risks of agentic AI according to the AI Exchange?" + ) + self.assertNotIn("Does the exchange cover", out) + system, user = llm.call_args.args + self.assertIn("Do not add, repeat or combine", system) + self.assertIn("Last message: What are some key risks", user) + + def test_model_failure_or_glued_questions_fall_back_to_the_current_message_only( + self, + ): + q = "What are some key risks" + self.assertEqual( + r.standalone_question( + q, HISTORY, mock.Mock(side_effect=RuntimeError("503")) + ), + q, + ) + glued = mock.Mock( + return_value="Does the exchange cover agentic AI? What are some key risks?" + ) + self.assertEqual(r.standalone_question(q, HISTORY, glued), q) + self.assertEqual( + r.standalone_question(q, HISTORY, mock.Mock(return_value="")), q + ) + + def test_no_history_returns_the_prompt(self): + self.assertEqual( + r.standalone_question("What are some key risks", None, mock.Mock()), + "What are some key risks", + ) + + def test_legacy_concat_mode_is_still_available(self): + with mock.patch.dict(os.environ, {"OWASP_HISTORY_MODE": "concat"}): + out = r.standalone_question( + "Berlin", [{"role": "user", "content": "events near me"}] + ) + self.assertEqual(out, "events near me\nBerlin") + + def test_effective_prompt_uses_the_standalone_logic(self): + q = "What are the three maturity levels defined by OWASP SAMM for each practice?" + self.assertEqual(r._effective_prompt(q, HISTORY), q) + + +class WebWarmUpTest(unittest.TestCase): + def test_warm_up_is_opt_in(self): + from application.utils.comprehension import shared + + app = mock.MagicMock() + with mock.patch.dict(os.environ, {"OWASP_WARMUP_WEB": ""}): + self.assertFalse(shared.start_background_warm_up(app)) + started = [] + with mock.patch.dict(os.environ, {"OWASP_WARMUP_WEB": "1"}), mock.patch( + "threading.Thread" + ) as thread, mock.patch.object(shared, "warm_up"): + thread.side_effect = ( + lambda target=None, **k: started.append(target) or mock.Mock() + ) + self.assertTrue(shared.start_background_warm_up(app)) + self.assertEqual(len(started), 1) + + +class ChatHandlerIsLazyTest(unittest.TestCase): + def test_agent_answers_do_not_build_the_cre_retrieval_handler(self): + from application import create_app, sqla + + env = { + "NO_LOGIN": "1", + "NO_LOAD_GRAPH_DB": "1", + "OWASP_AGENT_ENABLED": "1", + "INSECURE_REQUESTS": "1", + } + with mock.patch.dict(os.environ, env): + app = create_app(mode="test") + with app.app_context(): + sqla.create_all() + resp = { + "response": "Answer: x", + "accurate": True, + "owasp_agent": {"ok": True, "citations": []}, + } + with mock.patch( + "application.utils.owasp_agent.service.route_question", + return_value=resp, + ), mock.patch( + "application.prompt_client.prompt_client.PromptHandler" + ) as handler: + out = app.test_client().post( + "/rest/v1/completion", json={"prompt": "Who leads SAMM?"} + ) + sqla.session.remove() + sqla.drop_all() + self.assertEqual(out.status_code, 200) + handler.assert_not_called() + + +class TestCleanAnswerText(unittest.TestCase): + def test_links_and_double_brackets_become_plain_markers(self): + from application.utils.comprehension.synthesize import clean_answer_text + + raw = ( + "Owners attend [[7]](https://github.com/OWASP/x/blob/HEAD/a.md) and " + "[8](https://github.com/OWASP/x/blob/HEAD/b.md). See [the playbook](https://x.org/p) " + "or https://github.com/OWASP/y ." + ) + out = clean_answer_text(raw) + self.assertNotIn("http", out) + self.assertIn("[7]", out) + self.assertNotIn("[[7]]", out) + self.assertIn("[8].", out) + self.assertIn("See the playbook", out) + + def test_urls_containing_spaces_are_removed_too(self): + from application.utils.comprehension.synthesize import clean_answer_text + + raw = ( + "Include owners [[7]](https://github.com/OWASP/t/blob/HEAD/playbook/4. Train your " + "people to threat model.md) and testers [[8]](https://github.com/OWASP/t/a b.md)." + ) + self.assertEqual(clean_answer_text(raw), "Include owners [7] and testers [8].") + + def test_plain_markers_are_untouched(self): + from application.utils.comprehension.synthesize import clean_answer_text + + self.assertEqual( + clean_answer_text("SAMM has 3 levels [1][2]."), "SAMM has 3 levels [1][2]." + ) diff --git a/application/tests/chat_multiturn_test.py b/application/tests/chat_multiturn_test.py new file mode 100644 index 000000000..b168d8edf --- /dev/null +++ b/application/tests/chat_multiturn_test.py @@ -0,0 +1,240 @@ +"""Multi-turn chat: the conversation is background for a follow-up, never a second question. + +Each scripted conversation in data/multiturn_conversations.json is run through the real router with a +scripted rewrite model; we check what the answer pipeline actually receives. (scripts/ops/multiturn_probe.py +runs the same conversations against a live deployment.) +""" + +import json +import os +import tempfile +import unittest +from pathlib import Path +from unittest import mock + +from application.tests.owasp_agent.owasp_agent_test import _seed_store +from application.utils.owasp_agent.queries import MetaQueries +from application.utils.owasp_agent.router import OwaspAgentRouter + +CONVERSATIONS = json.loads( + (Path(__file__).parent / "data" / "multiturn_conversations.json").read_text() +) + + +class MultiTurnConversationsTest(unittest.TestCase): + def setUp(self): + self.tmp = tempfile.TemporaryDirectory() + store = _seed_store(os.path.join(self.tmp.name, "agent.sqlite")) + self.router = OwaspAgentRouter( + store=store, queries=MetaQueries(store, geocode_fn=lambda p: (0.0, 0.0)) + ) + env = mock.patch.dict( + os.environ, + { + "OWASP_AGENT_ENABLED": "1", + "OWASP_AGENT_SIM_ROUTER": "0", + "OWASP_HISTORY_MODE": "", + }, + ) + env.start() + self.addCleanup(env.stop) + self.addCleanup(self.tmp.cleanup) + + def _run(self, convo): + seen = [] + calls = [] + + def fake_rewriter(model, **kwargs): + def rewrite(system, user): + calls.append(user) + if convo["rewrite"] == "__FAIL__": + raise RuntimeError("503 model down") + return convo["rewrite"] or "" + + return rewrite + + def record(prompt): + seen.append(prompt) + return prompt + + with mock.patch( + "application.prompt_client.litellm_router.system_user_fn", + side_effect=fake_rewriter, + ), mock.patch( + "application.utils.owasp_agent.prompt_normalize.normalize_prompt", + side_effect=record, + ): + self.router.handle(convo["next"], history=convo["history"]) + return seen, calls + + def test_each_conversation_answers_only_its_own_question(self): + for convo in CONVERSATIONS: + with self.subTest(convo["name"]): + seen, calls = self._run(convo) + self.assertTrue(seen, "the answer pipeline never received the question") + received = seen[0] + self.assertEqual(received, convo["expect_standalone"]) + for forbidden in convo["must_not_contain"]: + self.assertNotIn(forbidden, received.lower()) + self.assertEqual(len(calls), convo["rewrite_calls"]) + + def test_the_rewrite_prompt_carries_the_history_as_background_and_the_last_message( + self, + ): + convo = CONVERSATIONS[0] + _, calls = self._run(convo) + self.assertIn("user: Does the exchange cover agentic AI?", calls[0]) + self.assertIn("assistant: Yes, the AI Exchange covers agentic AI", calls[0]) + self.assertIn("Last message: What are some key risks", calls[0]) + + def test_a_long_conversation_only_sends_the_recent_turns_to_the_rewrite(self): + history = [] + for i in range(10): + history += [ + {"role": "user", "content": f"question number {i}"}, + {"role": "assistant", "content": f"answer number {i}"}, + ] + convo = {**CONVERSATIONS[0], "history": history} + _, calls = self._run(convo) + self.assertNotIn("question number 0", calls[0]) + self.assertIn("answer number 9", calls[0]) + + +class ProbeEvaluationTest(unittest.TestCase): + def test_repeating_the_previous_answer_or_missing_the_topic_fails(self): + from scripts.ops.multiturn_probe import evaluate + + first = ( + "Answer: The AI Exchange covers agentic AI threats and controls in depth." + ) + self.assertEqual( + evaluate( + first, + "Key risks include prompt injection and excessive agency.", + ["risk"], + ), + [], + ) + self.assertTrue(evaluate(first, first, ["risk"])) # same text again + self.assertTrue(evaluate(first, "Something unrelated.", ["risk"])) + + +class NativeHistoryTest(unittest.TestCase): + def setUp(self): + from application.utils.comprehension import synthesize as syn + + self.syn = syn + self.addCleanup(syn.set_history, []) + + def test_history_becomes_real_turns_between_system_and_the_final_message(self): + self.syn.set_history( + [ + {"role": "user", "content": "Does the exchange cover agentic AI?"}, + {"role": "assistant", "content": "Yes."}, + ] + ) + out = self.syn.with_history( + [ + {"role": "system", "content": "S"}, + {"role": "user", "content": "Sources...\n\nQuestion: key risks?"}, + ] + ) + self.assertEqual( + [m["role"] for m in out], ["system", "user", "assistant", "user"] + ) + self.assertEqual(out[1]["content"], "Does the exchange cover agentic AI?") + self.assertIn("Question: key risks?", out[-1]["content"]) + + def test_history_is_trimmed_alternating_and_starts_with_a_user_turn(self): + long = "x" * 5000 + self.syn.set_history( + [{"role": "assistant", "content": "orphan"}] + + [ + {"role": "user", "content": long}, + {"role": "user", "content": "second user message"}, + ] + + [{"role": "assistant", "content": "a"}] * 2 + ) + turns = self.syn.get_history() + self.assertEqual( + [t["role"] for t in turns], ["user", "assistant"] + ) # orphan dropped, same roles merged + self.assertLessEqual( + len(turns[0]["content"]), self.syn.MAX_HISTORY_CHARS * 2 + 1 + ) + + def test_no_history_leaves_messages_alone_and_the_rule_is_only_added_with_history( + self, + ): + self.syn.set_history([]) + msgs = [{"role": "system", "content": "S"}, {"role": "user", "content": "Q"}] + self.assertEqual(self.syn.with_history(msgs), msgs) + seen = [] + self.syn.synthesize( + "q", [("body", ["u"])], lambda system, user: seen.append(system) or "A [1]" + ) + self.assertNotIn(self.syn.HISTORY_RULE, seen[0]) + self.syn.set_history([{"role": "user", "content": "earlier"}]) + self.syn.synthesize( + "q", [("body", ["u"])], lambda system, user: seen.append(system) or "A [1]" + ) + self.assertIn(self.syn.HISTORY_RULE, seen[1]) + + def test_router_hands_the_history_to_the_answer_call_once_per_turn(self): + import tempfile + + from application.tests.owasp_agent.owasp_agent_test import _seed_store + from application.utils.owasp_agent.queries import MetaQueries + from application.utils.owasp_agent.router import OwaspAgentRouter + + history = [ + {"role": "user", "content": "What is SAMM?"}, + {"role": "assistant", "content": "A maturity model."}, + ] + with tempfile.TemporaryDirectory() as tmp, mock.patch.dict( + os.environ, {"OWASP_AGENT_ENABLED": "1", "OWASP_AGENT_SIM_ROUTER": "0"} + ): + store = _seed_store(os.path.join(tmp, "a.sqlite")) + router = OwaspAgentRouter( + store=store, queries=MetaQueries(store, geocode_fn=lambda p: (0.0, 0.0)) + ) + router.handle("Who leads the Berlin chapter?", history=history) + self.assertEqual( + [t["role"] for t in self.syn.get_history()], ["user", "assistant"] + ) + + +class CitationNumberingTest(unittest.TestCase): + def test_markers_are_renumbered_to_match_the_returned_source_list(self): + from application.utils.comprehension import synthesize as syn + + hits = [ + ("s1", ["https://x/a"]), + ("s2", ["https://x/b"]), + ("s3", ["https://x/a"]), + ("s4", []), + ("s5", ["https://x/c"]), + ] + answer, cites = syn.synthesize( + "q", + hits, + lambda s, u: "Detection is imperfect [3][5]. Some attacks win [1][3][4]. Limits [2].", + ) + self.assertEqual(cites, ["https://x/a", "https://x/b", "https://x/c"]) + # [1] and [3] are two chunks of the same file -> one number; [4] has no URL -> dropped. + self.assertEqual( + answer, "Detection is imperfect [1][3]. Some attacks win [1]. Limits [2]." + ) + evidence = syn.take_evidence() + self.assertEqual(sorted({e["n"] for e in evidence if e["n"]}), [1, 2, 3]) + self.assertIn( + None, [e["n"] for e in evidence] + ) # the source without a link keeps its snippet, no number + + def test_answer_without_markers_is_untouched(self): + from application.utils.comprehension import synthesize as syn + + answer, cites = syn.synthesize( + "q", [("s1", ["https://x/a"])], lambda s, u: "No markers here." + ) + self.assertEqual((answer, cites), ("No markers here.", [])) diff --git a/application/tests/chat_unverified_test.py b/application/tests/chat_unverified_test.py new file mode 100644 index 000000000..f19287ee8 --- /dev/null +++ b/application/tests/chat_unverified_test.py @@ -0,0 +1,67 @@ +import os +import unittest +from unittest import mock + +from application import create_app, sqla +from application.utils.comprehension.answer import ABSTAIN_MESSAGE + + +class UnverifiedChatAnswerTest(unittest.TestCase): + def setUp(self): + env = mock.patch.dict( + os.environ, + { + "NO_LOGIN": "1", + "NO_LOAD_GRAPH_DB": "1", + "OWASP_AGENT_ENABLED": "0", + "INSECURE_REQUESTS": "1", + }, + ) + env.start() + self.addCleanup(env.stop) + self.app = create_app(mode="test") + self.ctx = self.app.app_context() + self.ctx.push() + sqla.create_all() + self.addCleanup(self._cleanup) + + def _cleanup(self): + sqla.session.remove() + sqla.drop_all() + self.ctx.pop() + + def _chat(self, generated): + handler = mock.MagicMock() + handler.generate_text.return_value = generated + with mock.patch( + "application.prompt_client.prompt_client.PromptHandler", + return_value=handler, + ): + return ( + self.app.test_client() + .post("/rest/v1/completion", json={"prompt": "who started opencre"}) + .get_json() + ) + + def test_a_free_form_answer_without_a_matching_source_is_replaced_by_an_honest_refusal( + self, + ): + out = self._chat( + { + "response": "OpenCRE was started by the Cloud Security Alliance.", + "table": "x", + "accurate": False, + } + ) + self.assertEqual(out["response"], ABSTAIN_MESSAGE) + self.assertNotIn("Cloud Security Alliance", str(out)) + self.assertFalse(out["accurate"]) + + def test_an_answer_backed_by_a_matching_node_is_kept(self): + out = self._chat({"response": "CRE 616-305 is about X.", "accurate": True}) + self.assertEqual(out["response"], "CRE 616-305 is about X.") + + def test_the_switch_restores_the_old_behaviour(self): + with mock.patch.dict(os.environ, {"CHAT_BLOCK_UNVERIFIED": "0"}): + out = self._chat({"response": "free text", "accurate": False}) + self.assertEqual(out["response"], "free text") diff --git a/application/tests/comprehension_authority_test.py b/application/tests/comprehension_authority_test.py new file mode 100644 index 000000000..23e6f8ac3 --- /dev/null +++ b/application/tests/comprehension_authority_test.py @@ -0,0 +1,75 @@ +import unittest +from types import SimpleNamespace as H + +from application.utils.comprehension import authority +from application.utils.comprehension.retrieve import _slug_of + + +def hit(slug, score, n=0): + return H(url=f"owasp-agent://evidence/comprehension-{slug}-raw-{n}", score=score) + + +class TestAuthority(unittest.TestCase): + def setUp(self): + self._orig = authority.weights + self.addCleanup(lambda: setattr(authority, "weights", self._orig)) + + def test_validate_drops_neutral_and_rejects_bad_values(self): + self.assertEqual( + authority.validate({"a": 1, "b": "1.2", "c": 0}), {"b": 1.2, "c": 0.0} + ) + for bad in ({"a": 5}, {"a": -1}, {"a": "x"}, {"a": True}, {"": 1.1}, [1]): + with self.assertRaises(ValueError): + authority.validate(bad) + + def test_bias_softens_and_hides_but_never_boosts(self): + authority.weights = lambda fresh=False: { + "aix": 1.3, + "noisy": 0.0, + "meh": 0.5, + } + hits = [ + hit("aisvs", 0.70), + hit("noisy", 0.69), + hit("meh", 0.68), + hit("aix", 0.60), + ] + out = authority.bias(hits, _slug_of) + self.assertEqual([_slug_of(h) for h in out], ["aisvs", "aix", "meh"]) + + def test_featured_tolerance_grows_with_weight(self): + self.assertAlmostEqual(authority.tolerance(1.15), 0.08) + self.assertGreater(authority.tolerance(1.3), authority.tolerance(1.15)) + self.assertEqual(authority.tolerance(5), 0.2) + self.assertEqual(authority.tolerance(1.01), 0.02) + + def test_featured_urls_roundtrip(self): + authority.reset_featured() + self.assertEqual(authority.featured_urls(), []) + authority.set_featured(["https://x"]) + self.assertEqual(authority.featured_urls(), ["https://x"]) + authority.reset_featured() + + def test_bias_is_a_noop_without_weights(self): + authority.weights = lambda fresh=False: {} + hits = [hit("a", 0.5), hit("b", 0.4)] + self.assertEqual(authority.bias(hits, _slug_of), hits) + + def test_diversify_gives_each_top_project_a_share(self): + hits = [hit("aisvs", 0.9, i) for i in range(8)] + [ + hit("aix", 0.5, 0), + hit("llm", 0.4, 0), + ] + out = authority.diversify(hits, _slug_of, projects=3, per=2) + self.assertEqual( + [_slug_of(h) for h in out[:4]], ["aisvs", "aisvs", "aix", "llm"] + ) + self.assertEqual(len(out), len(hits)) + + def test_diversify_leaves_single_project_results_alone(self): + hits = [hit("a", 0.9, i) for i in range(3)] + self.assertEqual(authority.diversify(hits, _slug_of), hits) + + +if __name__ == "__main__": + unittest.main() diff --git a/application/tests/comprehension_test.py b/application/tests/comprehension_test.py index f48eedd5f..f1d635682 100644 --- a/application/tests/comprehension_test.py +++ b/application/tests/comprehension_test.py @@ -1006,21 +1006,6 @@ def get(path): self.assertTrue(all("page_size=10" in c for c in calls)) -class TalksCatalogueTest(unittest.TestCase): - def test_parse_and_chunk(self): - from application.utils.comprehension import talks_source as ts - - talks = ts.parse_catalogue( - "abc|20241022|2667|Shift Left - Jane Doe\nbad line\ndef|20261021|60|Newer talk - John Roe\n" - ) - self.assertEqual(len(talks), 2) - chunks = ts.catalogue_chunks(talks, per_chunk=1) - self.assertEqual(len(chunks), 2) - self.assertIn("Newer talk - John Roe", chunks[0][2]) # newest first - self.assertIn("2024-10-22 (approximate)", chunks[1][2]) - self.assertIn("https://youtu.be/abc", chunks[1][2]) - - class NestBrokenPageTest(unittest.TestCase): def test_a_failing_page_is_read_item_by_item(self): from application.utils.comprehension import nest_source as ns @@ -1035,3 +1020,154 @@ def get(path): got = [i["k"] for i in ns.iter_items(get, "projects/")] self.assertEqual(got, [1, 3]) + + +class PickProjectTest(unittest.TestCase): + def _pick(self, reply): + from application.utils.comprehension import retrieve + + cands = [("owasp-asvs", "levels"), ("owasp-samm", "maturity")] + with mock.patch.object(retrieve, "pick_candidates", return_value=cands): + return retrieve.pick_project(None, "q", llm=lambda s, u: reply) + + def test_confident_choice_is_returned(self): + self.assertEqual(self._pick('{"project": 2, "confidence": 0.9}'), "owasp-samm") + + def test_unsure_or_malformed_means_no_pick(self): + self.assertEqual(self._pick('{"project": null, "confidence": 0.9}'), "") + self.assertEqual(self._pick('{"project": 1, "confidence": 0.4}'), "") + self.assertEqual(self._pick('{"project": 7, "confidence": 0.9}'), "") + self.assertEqual(self._pick("not json"), "") + + +class CompletePassTest(unittest.TestCase): + def test_missing_parts_parses_checker_reply(self): + from application.utils.comprehension.answer import missing_parts + + self.assertEqual( + missing_parts("q", "a", lambda s, u: '{"missing": ["the date", ""]}'), + ["the date"], + ) + self.assertEqual(missing_parts("q", "a", lambda s, u: "garbage"), []) + + def test_second_pass_reanswers_only_when_something_is_missing(self): + from application.utils.comprehension import answer as ans + + calls = [] + + def fake_answer(q, cands, complete): + calls.append(q) + return ("full answer", ["u"]) + + with mock.patch.object(ans, "_answer", side_effect=fake_answer), mock.patch( + "application.utils.comprehension.flags.helper_llm", + return_value=lambda s, u: '{"missing": ["x"]}', + ): + out = ans._complete_pass("q?", ("short", []), [], lambda m: "") + self.assertEqual(out[0], "full answer") + self.assertIn("must also cover: x", calls[0]) + with mock.patch.object(ans, "_answer", side_effect=fake_answer), mock.patch( + "application.utils.comprehension.flags.helper_llm", + return_value=lambda s, u: '{"missing": []}', + ): + calls.clear() + out = ans._complete_pass("q?", ("short", []), [], lambda m: "") + self.assertEqual(out[0], "short") + self.assertEqual(calls, []) + + +class NegCountRulesTest(unittest.TestCase): + def test_negation_detector(self): + from application.utils.comprehension.synthesize import NEG_RE + + for q in ( + "Which of these is NOT listed?", + "Everything except level 1", + "What isn't covered?", + ): + self.assertTrue(NEG_RE.search(q), q) + self.assertFalse(NEG_RE.search("How many levels does SAMM define?")) + + def test_whole_project_uses_count_and_neg_systems(self): + from application.utils.comprehension import answer as ans + from application.utils.comprehension.synthesize import COUNT_SYSTEM, NEG_RULE + + seen = [] + with mock.patch.dict( + os.environ, {"OWASP_CX_COUNT": "1", "OWASP_CX_NEG": "1"} + ), mock.patch( + "application.utils.comprehension.retrieve.project_sources", + return_value=[("X / a.md: text", ["u"])], + ), mock.patch( + "application.utils.comprehension.retrieve.named_slugs", + return_value=["owasp-x"], + ), mock.patch( + "application.utils.comprehension.answer.synthesize", + side_effect=lambda q, s, llm, system="": seen.append(system) + or ("a", ["u"]), + ): + ans._answer_from_whole_project( + "How many items are NOT listed?", lambda m: "" + ) + self.assertTrue(seen[0].startswith(COUNT_SYSTEM)) + self.assertTrue(seen[0].endswith(NEG_RULE)) + + +class AnswerUnavailableTest(unittest.TestCase): + def test_model_failure_is_not_reported_as_not_found(self): + from application.utils.comprehension import answer as ans + + with mock.patch.dict( + os.environ, {"OWASP_COMPREHENSION_ANSWER": "1"} + ), mock.patch.object( + ans, "_answer", side_effect=RuntimeError("503 service unavailable") + ), mock.patch.object( + ans, "enabled", return_value=True + ): + with self.assertRaises(ans.AnswerUnavailable): + ans.answer_from_candidates("q?", [], lambda m: "") + + +class PrefixCacheTest(unittest.TestCase): + def test_sources_come_first_when_flag_is_on(self): + from application.utils.comprehension.synthesize import synthesize + + seen = [] + + def llm(system, user): + seen.append(user) + return "A [1]" + + with mock.patch.dict(os.environ, {"OWASP_CX_PREFIXCACHE": "1"}): + synthesize("what?", [("body text", ["u"])], llm) + with mock.patch.dict(os.environ, {"OWASP_CX_PREFIXCACHE": ""}): + synthesize("what?", [("body text", ["u"])], llm) + self.assertTrue(seen[0].startswith("Sources:")) + self.assertTrue(seen[0].endswith("Question: what?")) + self.assertTrue(seen[1].startswith("Question: what?")) + + +class WikiChunksTest(unittest.TestCase): + def test_wiki_pages_become_chunks_with_a_wiki_path(self): + import tempfile + from pathlib import Path + + from application.utils.comprehension import sources + + with tempfile.TemporaryDirectory() as tmp: + d = Path(tmp) / "OWASP__x__wiki" + d.mkdir() + (d / "Home.md").write_text( + "# Home\n\n" + "The project wiki explains setup. " * 6 + ) + (d / "image.png").write_bytes(b"\x89PNG") + chunks = sources.wiki_chunks("OWASP/x", tmp) + self.assertTrue(chunks) + self.assertTrue(all(c[0] == "wiki/Home.md" for c in chunks)) + + +class WideMatchTest(unittest.TestCase): + def test_flag_registered(self): + from application.utils.comprehension import flags + + self.assertIn("WIDEMATCH", flags.NAMES) diff --git a/application/tests/cre_attach_test.py b/application/tests/cre_attach_test.py new file mode 100644 index 000000000..fc326f462 --- /dev/null +++ b/application/tests/cre_attach_test.py @@ -0,0 +1,92 @@ +import unittest +from unittest import mock + +from application.utils.owasp_agent import service + + +def _resp(kind="evidence", accurate=True, table=None): + return { + "response": "Answer: x", + "accurate": accurate, + "table": table if table is not None else [], + "owasp_agent": {"kind": kind, "citations": ["https://a/1"]}, + } + + +class TestAttachClosestCre(unittest.TestCase): + def test_evidence_answer_gets_the_closest_cre_first(self): + row = {"doctype": "CRE", "id": "123-456", "name": "Input validation"} + with mock.patch.object(service, "closest_cre_rows", return_value=[row]): + out = service.attach_closest_cre( + "q", _resp(table=[{"doctype": "Standard"}]) + ) + self.assertEqual(out["table"][0], row) + self.assertEqual(len(out["table"]), 2) + self.assertEqual( + out["owasp_agent"]["citations"][-1], "https://opencre.org/cre/123-456" + ) + self.assertEqual( + out["owasp_agent"]["citations"][0], "https://a/1" + ) # [n] markers stay valid + + def test_near_ties_are_all_listed_and_cited(self): + rows = [ + {"doctype": "CRE", "id": "366-835", "name": "Escape output against XSS"}, + {"doctype": "CRE", "id": "760-765", "name": "XSS protection"}, + ] + with mock.patch.object(service, "closest_cre_rows", return_value=rows): + out = service.attach_closest_cre("q", _resp()) + self.assertEqual([r["id"] for r in out["table"]], ["366-835", "760-765"]) + self.assertEqual( + out["owasp_agent"]["citations"][1:], + ["https://opencre.org/cre/366-835", "https://opencre.org/cre/760-765"], + ) + + def test_other_kinds_abstentions_and_no_match_are_untouched(self): + row = {"doctype": "CRE", "id": "1-2"} + with mock.patch.object(service, "closest_cre_rows", return_value=[row]): + self.assertEqual( + service.attach_closest_cre("q", _resp(kind="chapter_lookup"))["table"], + [], + ) + self.assertEqual( + service.attach_closest_cre("q", _resp(accurate=False))["table"], [] + ) + with mock.patch.object(service, "closest_cre_rows", return_value=[]): + self.assertEqual(service.attach_closest_cre("q", _resp())["table"], []) + + def test_lookup_failure_never_breaks_the_answer(self): + with mock.patch( + "application.utils.comprehension.shared.shared_handler", + side_effect=RuntimeError("boom"), + ): + self.assertEqual(service.closest_cre_rows("q"), []) + + +class TestSignificantWords(unittest.TestCase): + def test_question_words_that_name_a_cre_are_found(self): + q = service._significant_words("How do I prevent XSS?") + self.assertEqual(q, {"xss"}) + self.assertTrue(q & service._significant_words("XSS protection")) + self.assertTrue(q & service._significant_words("Escape output against XSS")) + self.assertFalse( + q & service._significant_words("Sanitize untrusted HTML input") + ) + + +class TestEncodeCitationUrls(unittest.TestCase): + def test_spaces_in_urls_are_encoded_and_positions_kept(self): + resp = {"owasp_agent": {"citations": ["https://x/a b.md", "https://y/ok"]}} + out = service.encode_citation_urls(resp) + self.assertEqual( + out["owasp_agent"]["citations"], ["https://x/a%20b.md", "https://y/ok"] + ) + + def test_clean_responses_are_returned_unchanged(self): + resp = {"owasp_agent": {"citations": ["https://y/ok"]}} + self.assertIs(service.encode_citation_urls(resp), resp) + self.assertEqual(service.encode_citation_urls("x"), "x") + + +if __name__ == "__main__": + unittest.main() diff --git a/application/tests/data/multiturn_conversations.json b/application/tests/data/multiturn_conversations.json new file mode 100644 index 000000000..e1406e717 --- /dev/null +++ b/application/tests/data/multiturn_conversations.json @@ -0,0 +1,85 @@ +[ + { + "name": "follow-up about risks must not re-answer the earlier question (reported 2026-10-24)", + "history": [ + {"role": "user", "content": "Does the exchange cover agentic AI?"}, + {"role": "assistant", "content": "Yes, the AI Exchange covers agentic AI threats and controls."} + ], + "next": "What are some key risks", + "rewrite": "What are some key risks of agentic AI according to the AI Exchange?", + "expect_standalone": "What are some key risks of agentic AI according to the AI Exchange?", + "must_not_contain": ["does the exchange cover"], + "rewrite_calls": 1 + }, + { + "name": "elliptical follow-up keeps the topic", + "history": [ + {"role": "user", "content": "What does OWASP membership cost in Switzerland?"}, + {"role": "assistant", "content": "Answer: Standard individual dues apply in Switzerland."} + ], + "next": "and in Morocco?", + "rewrite": "What does OWASP membership cost in Morocco?", + "expect_standalone": "What does OWASP membership cost in Morocco?", + "must_not_contain": ["switzerland"], + "rewrite_calls": 1 + }, + { + "name": "a self-contained question after an unrelated one ignores the history completely", + "history": [ + {"role": "user", "content": "What does ASVS cover?"}, + {"role": "assistant", "content": "ASVS is a verification standard."} + ], + "next": "Who are the project leaders of the OWASP Developer Guide project?", + "rewrite": null, + "expect_standalone": "Who are the project leaders of the OWASP Developer Guide project?", + "must_not_contain": ["asvs"], + "rewrite_calls": 0 + }, + { + "name": "third turn only depends on the previous answer, not on the first question", + "history": [ + {"role": "user", "content": "What is OWASP SAMM?"}, + {"role": "assistant", "content": "SAMM is the Software Assurance Maturity Model."}, + {"role": "user", "content": "How many levels does it define?"}, + {"role": "assistant", "content": "Three maturity levels per practice."} + ], + "next": "and what is level two?", + "rewrite": "What is maturity level two in OWASP SAMM?", + "expect_standalone": "What is maturity level two in OWASP SAMM?", + "must_not_contain": ["what is owasp samm?", "how many levels"], + "rewrite_calls": 1 + }, + { + "name": "the rewrite model glues the earlier question in: it is rejected, only the current message is used", + "history": [ + {"role": "user", "content": "Does the exchange cover agentic AI?"}, + {"role": "assistant", "content": "Yes."} + ], + "next": "What are some key risks", + "rewrite": "Does the exchange cover agentic AI? What are some key risks?", + "expect_standalone": "What are some key risks", + "must_not_contain": ["does the exchange cover"], + "rewrite_calls": 1 + }, + { + "name": "the rewrite model is down: current message alone, never the history", + "history": [ + {"role": "user", "content": "Does the exchange cover agentic AI?"}, + {"role": "assistant", "content": "Yes."} + ], + "next": "Why is that?", + "rewrite": "__FAIL__", + "expect_standalone": "Why is that?", + "must_not_contain": ["does the exchange cover"], + "rewrite_calls": 1 + }, + { + "name": "first message of a chat has no history", + "history": [], + "next": "What are some key risks", + "rewrite": null, + "expect_standalone": "What are some key risks", + "must_not_contain": [], + "rewrite_calls": 0 + } +] diff --git a/application/tests/llm_resilience_test.py b/application/tests/llm_resilience_test.py new file mode 100644 index 000000000..4dc67bcae --- /dev/null +++ b/application/tests/llm_resilience_test.py @@ -0,0 +1,142 @@ +import os +import unittest +from unittest import mock + +from application.prompt_client import litellm_router as lr +from application.prompt_client.llm_error_utils import ( + is_overload_error, + is_transient_llm_error, +) + + +class ServiceUnavailableError(Exception): + pass + + +class AuthenticationError(Exception): + pass + + +def ok(text="hi"): + msg = mock.Mock(content=text) + return mock.Mock(choices=[mock.Mock(message=msg)]) + + +class ClassifierTest(unittest.TestCase): + def test_provider_outages_are_transient(self): + for err in ( + ServiceUnavailableError( + "GeminiException - The model is overloaded. high demand" + ), + RuntimeError("503 Service Unavailable"), + RuntimeError('{"error": {"code": 503, "message": "x"}}'), + RuntimeError("Internal error encountered"), + TimeoutError("timed out"), + ): + self.assertTrue(is_transient_llm_error(err), err) + self.assertTrue(is_overload_error(ServiceUnavailableError("x"))) + + def test_bad_key_or_request_is_not_transient(self): + self.assertFalse( + is_transient_llm_error(AuthenticationError("API key not valid")) + ) + self.assertFalse(is_transient_llm_error(ValueError("bad request"))) + + +class ResilientCompletionTest(unittest.TestCase): + def setUp(self): + lr.reset_breaker() + env = mock.patch.dict( + os.environ, + { + "LLM_INTERACTIVE_RETRIES": "2", + "LLM_INTERACTIVE_BASE_DELAY": "0.01", + "LLM_BREAKER_FAILURES": "2", + "LLM_BREAKER_COOLDOWN_SECONDS": "60", + }, + ) + env.start() + self.addCleanup(env.stop) + sleeper = mock.patch("time.sleep") + self.sleep = sleeper.start() + self.addCleanup(sleeper.stop) + + def _client(self, behaviour): + client = mock.Mock() + client.completion.side_effect = behaviour + return client + + def test_recovers_from_a_short_outage_without_fallback(self): + calls = [] + + def behaviour(model, messages, **kw): + calls.append(model) + if len(calls) < 3: + raise ServiceUnavailableError("503 overloaded") + return ok("answer") + + out = lr.resilient_completion( + model="m1", + fallback_model="m2", + messages=[{"role": "user", "content": "q"}], + client=self._client(behaviour), + ) + self.assertEqual(lr.extract_content_text(out), "answer") + self.assertEqual( + calls, ["m1", "m1", "m1"] + ) # same model after two quick retries + + def test_falls_back_when_the_primary_stays_down(self): + def behaviour(model, messages, **kw): + if model == "m1": + raise ServiceUnavailableError("503 overloaded") + return ok("from fallback") + + out = lr.resilient_completion( + model="m1", fallback_model="m2", messages=[], client=self._client(behaviour) + ) + self.assertEqual(lr.extract_content_text(out), "from fallback") + + def test_non_transient_errors_do_not_retry_or_fall_back(self): + client = self._client(AuthenticationError("API key not valid")) + with self.assertRaises(AuthenticationError): + lr.resilient_completion( + model="m1", fallback_model="m2", messages=[], client=client + ) + self.assertEqual(client.completion.call_count, 1) + + def test_breaker_skips_the_failing_primary_after_repeated_failures(self): + seen = [] + + def behaviour(model, messages, **kw): + seen.append(model) + if model == "m1": + raise ServiceUnavailableError("503") + return ok("fb") + + client = self._client(behaviour) + for _ in range(2): # two full failures open the breaker (threshold 2) + lr.resilient_completion( + model="m1", fallback_model="m2", messages=[], client=client + ) + seen.clear() + lr.resilient_completion( + model="m1", fallback_model="m2", messages=[], client=client + ) + self.assertEqual(seen, ["m2"]) # straight to the fallback + + def test_without_a_fallback_an_outage_is_raised(self): + client = self._client(ServiceUnavailableError("503")) + with self.assertRaises(ServiceUnavailableError): + lr.resilient_completion(model="m1", messages=[], client=client) + + def test_quick_retry_helper_for_embeddings(self): + n = {"i": 0} + + def flaky(): + n["i"] += 1 + if n["i"] < 3: + raise RuntimeError("503 Service Unavailable") + return [0.1] + + self.assertEqual(lr.quick_retry(flaky, context="t"), [0.1]) diff --git a/application/tests/mcp_ask_test.py b/application/tests/mcp_ask_test.py index e20854277..b7caa262a 100644 --- a/application/tests/mcp_ask_test.py +++ b/application/tests/mcp_ask_test.py @@ -106,7 +106,7 @@ def test_server_always_lists_ask(self) -> None: instructions = server_instructions() self.assertIn("ask tool", instructions.lower()) - @patch("application.mcp.ask.OwaspAgentRouter") + @patch("application.utils.owasp_agent.OwaspAgentRouter") def test_auto_meta_channel(self, router_cls: MagicMock) -> None: router_cls.return_value.handle.return_value = { "response": "Answer: 1 chapter", @@ -130,7 +130,7 @@ def test_auto_meta_channel(self, router_cls: MagicMock) -> None: self.assertIn("https://nest.example/la", out["citations"]) router_cls.return_value.handle.assert_called_once() - @patch("application.mcp.ask.OwaspAgentRouter") + @patch("application.utils.owasp_agent.OwaspAgentRouter") def test_meta_runs_before_cre_normative_defer(self, router_cls: MagicMock) -> None: """AIX/meta evidence must win even when intent is cre_normative.""" router_cls.return_value.handle.return_value = { @@ -232,7 +232,7 @@ def test_corpus_security_skips_router(self, retrieve: MagicMock) -> None: } }, ) - with patch("application.mcp.ask.OwaspAgentRouter") as router_cls: + with patch("application.utils.owasp_agent.OwaspAgentRouter") as router_cls: out = ask( "authentication", corpus=CORPUS_SECURITY, @@ -290,7 +290,9 @@ def test_corpus_owasp_meta_rejects_cre_looking(self) -> None: with patch( "application.mcp.ask.classify_intent", return_value="cre_normative" ): - with patch("application.mcp.ask.OwaspAgentRouter") as router_cls: + with patch( + "application.utils.owasp_agent.OwaspAgentRouter" + ) as router_cls: router_cls.return_value.handle.return_value = None out = ask("How should I store passwords?", corpus=CORPUS_OWASP_META) self.assertFalse(out["ok"]) diff --git a/application/tests/owasp_agent/node_facts_test.py b/application/tests/owasp_agent/node_facts_test.py index e5afa058c..255665bba 100644 --- a/application/tests/owasp_agent/node_facts_test.py +++ b/application/tests/owasp_agent/node_facts_test.py @@ -270,7 +270,7 @@ class _Dec: self.assertNotIn(bad, out_msg) queries.project_lookup.assert_not_called() - @patch("application.mcp.ask.OwaspAgentRouter") + @patch("application.utils.owasp_agent.OwaspAgentRouter") def test_ask_auto_uses_fixture_prompts(self, router_cls: MagicMock) -> None: from application.mcp.ask import CORPUS_AUTO, ask diff --git a/application/tests/owasp_agent/owasp_agent_test.py b/application/tests/owasp_agent/owasp_agent_test.py index 332bc767d..da71e2bbb 100644 --- a/application/tests/owasp_agent/owasp_agent_test.py +++ b/application/tests/owasp_agent/owasp_agent_test.py @@ -958,10 +958,18 @@ def test_clarify_followup_uses_history(self) -> None: "content": "Answer: Standard individual dues apply in Switzerland.", }, ] - resp = self.router.handle("and in Morocco?", history=history) + # The follow-up is rewritten into a standalone question using the history as background. + rewrite = ( + lambda system, user: "What does OWASP membership cost in Morocco?" + ) # noqa: E731 + with patch( + "application.prompt_client.litellm_router.system_user_fn", + return_value=rewrite, + ): + resp = self.router.handle("and in Morocco?", history=history) self.assertIsNotNone(resp) assert resp is not None - # Combined context should route as membership, not bare clarify miss. + # The standalone question should route as membership, not bare clarify miss. self.assertIn("owasp_agent", resp) kind = resp["owasp_agent"]["kind"] self.assertIn(kind, ("membership", "clarify")) diff --git a/application/tests/rate_limit_test.py b/application/tests/rate_limit_test.py new file mode 100644 index 000000000..c23b212fb --- /dev/null +++ b/application/tests/rate_limit_test.py @@ -0,0 +1,75 @@ +import os +import unittest +from unittest import mock + +from application.utils import rate_limit, roles + + +class RateLimitTest(unittest.TestCase): + def setUp(self): + rate_limit.reset_memory() + self._env = mock.patch.dict( + os.environ, + { + "RATE_LIMIT_ENABLED": "1", + "RATE_LIMIT_PER_MINUTE": "3", + "RATE_LIMIT_PER_HOUR": "5", + "RATE_LIMIT_PER_DAY": "9", + "RATE_LIMIT_ADMIN_MULTIPLIER": "2", + }, + ) + self._env.start() + self.addCleanup(self._env.stop) + + def test_minute_limit_then_blocked_with_retry_after(self): + for _ in range(3): + rate_limit.check_and_count("a@owasp.org", roles.USER) + with self.assertRaises(rate_limit.RateLimited) as cm: + rate_limit.check_and_count("a@owasp.org", roles.USER) + self.assertEqual(cm.exception.window, "minute") + self.assertTrue(1 <= cm.exception.retry_after <= 60) + self.assertIn("3 questions per minute", cm.exception.message()) + + def test_limits_are_per_account_and_case_insensitive(self): + for _ in range(3): + rate_limit.check_and_count("A@owasp.org", roles.USER) + rate_limit.check_and_count( + "b@owasp.org", roles.USER + ) # a different account is unaffected + with self.assertRaises(rate_limit.RateLimited): + rate_limit.check_and_count("a@OWASP.org", roles.USER) + + def test_admins_get_a_multiple(self): + for _ in range(6): + rate_limit.check_and_count("boss@owasp.org", roles.ADMIN) + with self.assertRaises(rate_limit.RateLimited): + rate_limit.check_and_count("boss@owasp.org", roles.ADMIN) + + def test_disabled_and_anonymous_are_not_limited(self): + with mock.patch.dict(os.environ, {"RATE_LIMIT_ENABLED": "0"}): + for _ in range(50): + rate_limit.check_and_count("a@owasp.org", roles.USER) + for _ in range(50): + rate_limit.check_and_count("", roles.USER) + + def test_exempt_accounts_are_not_limited(self): + with mock.patch.dict( + os.environ, {"RATE_LIMIT_EXEMPT_EMAILS": "Bot@owasp.org, other@x.org"} + ): + for _ in range(50): + rate_limit.check_and_count("bot@owasp.org", roles.USER) + for _ in range(3): + rate_limit.check_and_count("jane@owasp.org", roles.USER) + with self.assertRaises(rate_limit.RateLimited): + rate_limit.check_and_count("jane@owasp.org", roles.USER) + + def test_hour_window_counts_across_minutes(self): + now = [10_000.0] + with mock.patch("time.time", side_effect=lambda: now[0]): + for i in range(5): + now[0] = 10_000.0 + i * 100 # a new minute each time, same hour + rate_limit.check_and_count("c@owasp.org", roles.USER) + now[0] = 10_000.0 + 5 * 100 + with self.assertRaises(rate_limit.RateLimited) as cm: + rate_limit.check_and_count("c@owasp.org", roles.USER) + self.assertEqual(cm.exception.window, "hour") diff --git a/application/tests/roles_access_test.py b/application/tests/roles_access_test.py new file mode 100644 index 000000000..98c0338dc --- /dev/null +++ b/application/tests/roles_access_test.py @@ -0,0 +1,264 @@ +"""Two-role access: any allowed-domain account can chat and use MCP ask; only named admins can do admin things.""" + +from __future__ import annotations + +import json +import os +import unittest +from unittest.mock import MagicMock, patch + +import mcp.types as types + +from application import create_app, sqla +from application.database import db +from application.tests.mcp_remote_http_test import McpRemoteHttpTest +from application.utils.api_tokens import mint as api_mint + +ROLE_ENV = { + "ADMIN_EMAILS": "boss@owasp.org", + "LOGIN_ALLOWED_DOMAINS": "owasp.org", + "LOGIN_ALLOWED_EMAILS": "", + "NO_LOGIN": "0", + "CRE_ENABLE_LOGIN": "1", + "INSECURE_REQUESTS": "1", + "NO_LOAD_GRAPH_DB": "1", + "FLASK_CONFIG": "production", +} + + +class WebRolesTest(unittest.TestCase): + def setUp(self) -> None: + self._env = patch.dict(os.environ, ROLE_ENV) + self._env.start() + self.addCleanup(self._env.stop) + self.app = create_app(mode="test") + self.app.secret_key = "test-secret" + self.ctx = self.app.app_context() + self.ctx.push() + sqla.create_all() + + def tearDown(self) -> None: + sqla.session.remove() + sqla.drop_all() + self.ctx.pop() + + def _client(self, email): + c = self.app.test_client() + with c.session_transaction() as s: + s["user_id"] = "u-" + email + s["email"] = email + s["google_id"] = "g-" + email + return c + + def test_me_reports_the_role(self): + self.assertEqual( + self._client("boss@owasp.org").get("/rest/v1/me").get_json()["role"], + "admin", + ) + self.assertEqual( + self._client("jane@owasp.org").get("/rest/v1/me").get_json()["role"], "user" + ) + + def test_account_outside_the_domain_is_refused_and_signed_out(self): + c = self._client("eve@gmail.com") + self.assertEqual(c.get("/rest/v1/me").status_code, 403) + with c.session_transaction() as s: + self.assertNotIn("user_id", s) + + def test_user_can_use_chat_endpoint_gate_but_not_admin(self): + c = self._client("jane@owasp.org") + handler = MagicMock() + handler.generate_text.return_value = {"response": "ok", "accurate": True} + with patch.dict(os.environ, {"OWASP_AGENT_ENABLED": "0"}), patch( + "application.prompt_client.prompt_client.PromptHandler", + return_value=handler, + ): + r = c.post("/rest/v1/completion", json={"prompt": "hi"}) + self.assertEqual(r.status_code, 200, r.data) + self.assertEqual(r.get_json()["response"], "ok") + self.assertEqual( + c.get("/rest/v1/api-keys").status_code, 403 + ) # admin-wide key list + self.assertEqual(c.get("/admin/config").status_code, 403) + with c.session_transaction() as s: + self.assertIn( + "user_id", s + ) # a plain user is NOT signed out for poking an admin URL + + def test_chat_is_rate_limited_per_account(self): + from application.utils import rate_limit + + rate_limit.reset_memory() + handler = MagicMock() + handler.generate_text.return_value = {"response": "ok", "accurate": True} + with patch.dict( + os.environ, + { + "OWASP_AGENT_ENABLED": "0", + "RATE_LIMIT_PER_MINUTE": "2", + "RATE_LIMIT_ENABLED": "1", + }, + ), patch( + "application.prompt_client.prompt_client.PromptHandler", + return_value=handler, + ): + jane, bob = self._client("jane@owasp.org"), self._client("bob@owasp.org") + codes = [ + jane.post("/rest/v1/completion", json={"prompt": "hi"}).status_code + for _ in range(3) + ] + third = jane.post("/rest/v1/completion", json={"prompt": "hi"}) + other = bob.post("/rest/v1/completion", json={"prompt": "hi"}).status_code + self.assertEqual(codes, [200, 200, 429]) + self.assertIn("Retry-After", third.headers) + self.assertIn("per minute", third.get_json()["error"]) + self.assertEqual(other, 200) # another account still has its own allowance + + def test_admin_passes_the_admin_gates(self): + c = self._client("boss@owasp.org") + self.assertEqual(c.get("/rest/v1/api-keys").status_code, 200) + + def test_user_manages_only_their_own_keys(self): + users = db.Node_collection() + jane = users.upsert_user( + google_sub="g-jane@owasp.org", email="jane@owasp.org", display_name="Jane" + ) + bob = users.upsert_user( + google_sub="g-bob@owasp.org", email="bob@owasp.org", display_name="Bob" + ) + bobs_key = api_mint.mint_api_token(user=bob, name="bob-key", expires_days=5) + c = self.app.test_client() + with c.session_transaction() as s: + s["user_id"], s["email"], s["google_id"], s["csrf_token"] = ( + jane.id, + "jane@owasp.org", + "g-jane@owasp.org", + "tok", + ) + hdr = {"X-CSRF-Token": "tok"} + made = c.post("/rest/v1/my-api-keys", json={"name": "mine"}, headers=hdr) + self.assertEqual(made.status_code, 201, made.data) + mine = c.get("/rest/v1/my-api-keys").get_json()["keys"] + self.assertEqual([k["name"] for k in mine], ["mine"]) + self.assertEqual( + c.delete(f"/rest/v1/my-api-keys/{bobs_key.jti}", headers=hdr).status_code, + 404, + ) + self.assertEqual( + c.delete( + f"/rest/v1/my-api-keys/{made.get_json()['jti']}", headers=hdr + ).status_code, + 204, + ) + + +class McpRolesTest(unittest.TestCase): + """Remote MCP with the two-role environment (borrows the harness helpers).""" + + _mint = McpRemoteHttpTest._mint + _call_tool = McpRemoteHttpTest._call_tool + + def setUp(self) -> None: + self._env = patch.dict( + os.environ, + { + **ROLE_ENV, + "API_TOKEN_JWT_SECRET": "test-api-token-jwt-secret-not-for-prod", + "OPENCRE_MCP_REMOTE": "1", + "OWASP_AGENT_ENABLED": "1", + }, + ) + self._env.start() + self.app = create_app(mode="test") + self.app.secret_key = "test-secret" + self.app_context = self.app.app_context() + self.app_context.push() + sqla.create_all() + self.collection = db.Node_collection() + self.boss = self.collection.upsert_user( + google_sub="boss-sub", email="boss@owasp.org", display_name="Boss" + ) + self.jane = self.collection.upsert_user( + google_sub="jane-sub", email="jane@owasp.org", display_name="Jane" + ) + self.eve = self.collection.upsert_user( + google_sub="eve-sub", email="eve@gmail.com", display_name="Eve" + ) + + def tearDown(self) -> None: + sqla.session.remove() + sqla.drop_all() + self.app_context.pop() + self._env.stop() + + def _text(self, result): + first = result.content[0] + assert isinstance(first, types.TextContent) + return first.text + + def test_user_may_ask_and_rate(self): + tok = self._mint(self.jane) + with patch("application.mcp.server.run_ask", return_value={"ok": True}) as ask: + r = self._call_tool( + "ask", {"prompt": "What is MFA?"}, authorization=f"Bearer {tok}" + ) + ask.assert_called_once() + self.assertFalse(r.is_error) + + def test_outsider_and_anonymous_are_refused(self): + with patch("application.mcp.server.run_ask") as ask: + eve = self._call_tool( + "ask", {"prompt": "x"}, authorization=f"Bearer {self._mint(self.eve)}" + ) + anon = self._call_tool("ask", {"prompt": "x"}) + ask.assert_not_called() + self.assertTrue(eve.is_error and anon.is_error) + self.assertRegex(self._text(eve).lower(), r"403|no access") + + def test_meta_and_source_tools_need_a_signed_in_role(self): + self.assertTrue( + self._call_tool("owasp_meta_sources_list", {"project": "asvs"}).is_error + ) # anonymous + with patch( + "application.mcp.server.call_owasp_meta_tool", return_value={"ok": True} + ): + r = self._call_tool( + "owasp_meta_sources_list", + {"project": "asvs"}, + authorization=f"Bearer {self._mint(self.jane)}", + ) + self.assertFalse(r.is_error, self._text(r)) + + def test_public_cre_read_tools_stay_public(self): + self.assertFalse(self._call_tool("list_standards", {}).is_error) + + def test_admin_only_tools_need_admin(self): + from application.mcp import remote_auth + + with patch.object(remote_auth, "ADMIN_ONLY_TOOLS", frozenset({"ask"})), patch( + "application.mcp.server.run_ask", return_value={"ok": True} + ): + user = self._call_tool( + "ask", {"prompt": "x"}, authorization=f"Bearer {self._mint(self.jane)}" + ) + admin = self._call_tool( + "ask", {"prompt": "x"}, authorization=f"Bearer {self._mint(self.boss)}" + ) + self.assertTrue(user.is_error) + self.assertRegex(self._text(user).lower(), r"admin role is required") + self.assertFalse(admin.is_error) + + def test_mcp_ask_is_rate_limited_per_account(self): + from application.utils import rate_limit + + rate_limit.reset_memory() + tok = self._mint(self.jane) + with patch.dict( + os.environ, {"RATE_LIMIT_PER_MINUTE": "2", "RATE_LIMIT_ENABLED": "1"} + ), patch("application.mcp.server.run_ask", return_value={"ok": True}): + results = [ + self._call_tool("ask", {"prompt": "x"}, authorization=f"Bearer {tok}") + for _ in range(3) + ] + self.assertEqual([r.is_error for r in results], [False, False, True]) + self.assertTrue(self._text(results[2]).startswith("429 ")) diff --git a/application/tests/roles_test.py b/application/tests/roles_test.py new file mode 100644 index 000000000..cc44eba77 --- /dev/null +++ b/application/tests/roles_test.py @@ -0,0 +1,80 @@ +import os +import unittest +from unittest import mock + +from application.utils import roles + + +def env(**kw): + base = { + "ADMIN_EMAILS": "", + "LOGIN_ALLOWED_EMAILS": "", + "LOGIN_ALLOWED_DOMAINS": "", + "FLASK_CONFIG": "", + } + base.update(kw) + return mock.patch.dict(os.environ, base) + + +class RolesTest(unittest.TestCase): + def test_roles_model(self): + with env( + ADMIN_EMAILS="Boss@owasp.org, ops@example.com", + LOGIN_ALLOWED_DOMAINS="owasp.org", + ): + self.assertEqual(roles.role_for("boss@owasp.org"), roles.ADMIN) + self.assertEqual( + roles.role_for("ops@example.com"), roles.ADMIN + ) # named admin outside the domain + self.assertEqual(roles.role_for("anyone@owasp.org"), roles.USER) + self.assertIsNone(roles.role_for("x@gmail.com")) + self.assertIsNone( + roles.role_for("x@evilowasp.org") + ) # suffix tricks do not match + self.assertIsNone(roles.role_for("")) + self.assertTrue(roles.has_role("anyone@owasp.org", roles.USER)) + self.assertFalse(roles.has_role("anyone@owasp.org", roles.ADMIN)) + self.assertTrue(roles.has_role("boss@owasp.org", roles.USER)) + + def test_subdomain_and_star_domains(self): + with env(ADMIN_EMAILS="a@x.org", LOGIN_ALLOWED_DOMAINS="@owasp.org"): + self.assertEqual(roles.role_for("u@chapters.owasp.org"), roles.USER) + with env(ADMIN_EMAILS="a@x.org", LOGIN_ALLOWED_DOMAINS="*"): + self.assertEqual(roles.role_for("u@anything.io"), roles.USER) + + def test_no_domain_configured_fails_closed(self): + with env(ADMIN_EMAILS="a@x.org"): + self.assertIsNone(roles.role_for("u@owasp.org")) + self.assertEqual(roles.role_for("a@x.org"), roles.ADMIN) + + def test_legacy_mode_keeps_the_old_allowlist(self): + with env(LOGIN_ALLOWED_EMAILS="a@x.org", LOGIN_ALLOWED_DOMAINS="owasp.org"): + self.assertFalse(roles.roles_enabled()) + self.assertEqual(roles.role_for("a@x.org"), roles.ADMIN) + self.assertIsNone(roles.role_for("u@owasp.org")) + with env(FLASK_CONFIG="production"): + self.assertIsNone(roles.role_for("a@x.org")) + with env(): + self.assertEqual( + roles.role_for("a@x.org"), roles.ADMIN + ) # dev convenience as before + + +class McpToolRolesTest(unittest.TestCase): + def test_tool_roles(self): + from application.mcp import remote_auth as ra + + cre = frozenset({"list_standards"}) + self.assertEqual(ra.min_role_for_tool("ask", cre), roles.USER) + self.assertEqual(ra.min_role_for_tool("rate_answer", cre), roles.USER) + self.assertEqual( + ra.min_role_for_tool("owasp_meta_sources_search", cre), roles.USER + ) + self.assertIsNone( + ra.min_role_for_tool("list_standards", cre) + ) # public CRE read tool + self.assertEqual( + ra.min_role_for_tool("some_future_tool", cre), roles.ADMIN + ) # deny by default + with mock.patch.object(ra, "ADMIN_ONLY_TOOLS", frozenset({"ask"})): + self.assertEqual(ra.min_role_for_tool("ask", cre), roles.ADMIN) diff --git a/application/utils/answer_feedback.py b/application/utils/answer_feedback.py new file mode 100644 index 000000000..d3ffbb497 --- /dev/null +++ b/application/utils/answer_feedback.py @@ -0,0 +1,158 @@ +"""Log every answer and record the asker's thumbs up/down. + +Logging is best effort: a failure here must never break answering a question. +""" + +from __future__ import annotations + +import json +import os +from datetime import datetime, timezone +from typing import Any, Dict, List, Optional, Sequence + +from cre_logging import get_logger + +logger = get_logger(__name__) + +RATINGS = ("up", "down") +MAX_QUESTION = 4000 +MAX_ANSWER = 20000 +MAX_COMMENT = 2000 + + +def _utcnow() -> datetime: + return datetime.now(timezone.utc).replace(tzinfo=None) + + +def _truthy(name: str, default: str) -> bool: + return (os.environ.get(name, default) or default).strip().lower() not in ( + "0", + "false", + "no", + "off", + "", + ) + + +def logging_enabled(source: str) -> bool: + """MCP answers are logged by default. Chat answers only when a deployment opts in + (``ANSWER_LOG_CHAT=1``): the public chat page promises not to store questions.""" + if not _truthy("ANSWER_LOG_ENABLED", "1"): + return False + return _truthy("ANSWER_LOG_CHAT", "0") if source == "chat" else True + + +def _active_flags() -> str: + from application.utils.comprehension import flags + + return ",".join(flags.active()) + + +def log_answer( + *, + source: str, + question: str, + answer: Optional[str], + citations: Optional[Sequence[str]] = None, + channel: Optional[str] = None, + latency_ms: Optional[int] = None, +) -> Optional[str]: + """Store an answer and return its id (None when logging is off or fails).""" + if not logging_enabled(source): + return None + from application import sqla + from application.database import db + + try: + row = db.AnswerLog( + created_at=_utcnow(), + source=source, + question=(question or "")[:MAX_QUESTION], + answer=(answer or "")[:MAX_ANSWER], + citations=json.dumps(list(citations or [])), + channel=channel, + flags=_active_flags(), + latency_ms=latency_ms, + ) + sqla.session.add(row) + sqla.session.commit() + return str(row.id) + except Exception: # noqa: BLE001 -- never break answering + sqla.session.rollback() + logger.warning("answer log failed", exc_info=True) + return None + + +def record_rating( + answer_id: str, rating: str, comment: Optional[str] = None +) -> Dict[str, Any]: + """Set (or change) the thumbs for a logged answer. Raises ValueError / KeyError.""" + from application import sqla + from application.database import db + + rating = (rating or "").strip().lower() + if rating not in RATINGS: + raise ValueError(f"rating must be one of: {', '.join(RATINGS)}") + if not answer_id or not isinstance(answer_id, str): + raise ValueError("answer_id is required") + row = sqla.session.get(db.AnswerLog, answer_id) + if row is None: + raise KeyError("unknown answer_id") + row.rating = rating + row.comment = (comment or "").strip()[:MAX_COMMENT] or None + row.rated_at = _utcnow() + sqla.session.commit() + return {"answer_id": answer_id, "rating": rating} + + +def list_feedback( + rating: Optional[str] = None, limit: int = 100 +) -> List[Dict[str, Any]]: + """Newest rated answers first, for review and for building eval sets.""" + from application import sqla + from application.database import db + + query = sqla.session.query(db.AnswerLog).filter(db.AnswerLog.rating.isnot(None)) + if rating: + query = query.filter(db.AnswerLog.rating == rating) + rows = ( + query.order_by(db.AnswerLog.rated_at.desc()) + .limit(max(1, min(int(limit), 500))) + .all() + ) + return [ + { + "id": r.id, + "source": r.source, + "question": r.question, + "answer": r.answer, + "citations": json.loads(r.citations or "[]"), + "rating": r.rating, + "comment": r.comment, + "flags": r.flags, + "latency_ms": r.latency_ms, + "rated_at": r.rated_at.isoformat() if r.rated_at else None, + } + for r in rows + ] + + +def attach_answer_id( + result: Dict[str, Any], + *, + question: str, + source: str, + latency_ms: Optional[int] = None, +) -> Dict[str, Any]: + """Log an ``ask`` result and put the id on it so the client can rate it.""" + answer_id = log_answer( + source=source, + question=question, + answer=str(result.get("message") or ""), + citations=result.get("citations") or [], + channel=str(result.get("channel") or ""), + latency_ms=latency_ms, + ) + if answer_id: + result["answer_id"] = answer_id + return result diff --git a/application/utils/api_tokens/mint.py b/application/utils/api_tokens/mint.py index 43c25ad4b..a0d3f4963 100644 --- a/application/utils/api_tokens/mint.py +++ b/application/utils/api_tokens/mint.py @@ -90,8 +90,12 @@ def mint_api_token( ) -def revoke_api_token(jti: str) -> bool: - row = sqla.session.query(db.ApiKey).filter(db.ApiKey.jti == jti).first() +def revoke_api_token(jti: str, *, user_id: Optional[str] = None) -> bool: + """Revoke a key; with ``user_id`` only that user's own key can be revoked.""" + q = sqla.session.query(db.ApiKey).filter(db.ApiKey.jti == jti) + if user_id: + q = q.filter(db.ApiKey.user_id == user_id) + row = q.first() if row is None: return False if row.revoked_at is None: diff --git a/application/utils/api_tokens/verify.py b/application/utils/api_tokens/verify.py index ea430481e..082f2e6fe 100644 --- a/application/utils/api_tokens/verify.py +++ b/application/utils/api_tokens/verify.py @@ -85,22 +85,18 @@ def verify_api_token(raw: str) -> Optional[ApiKeyIdentity]: def email_on_login_allowlist(email: Optional[str]) -> bool: - """True when ``email`` is on LOGIN_ALLOWED_EMAILS. + """True when ``email`` has the admin role (kept under its old name for callers). - Mirrors Admin ``login_required`` / ``_login_email_allowed`` without Flask - request path gates. Dev/test: empty allowlist allows anyone. - ``FLASK_CONFIG=production``: empty allowlist fails closed. + Legacy mode: on LOGIN_ALLOWED_EMAILS (dev/test: empty allows anyone; production: empty + fails closed). Roles mode: one of the named ``ADMIN_EMAILS``. """ - raw = (os.environ.get("LOGIN_ALLOWED_EMAILS") or "").strip() - is_prod = (os.environ.get("FLASK_CONFIG") or "").strip().lower() == "production" - if not raw: - return not is_prod - allowed = {e.strip().lower() for e in raw.split(",") if e.strip()} - return (email or "").strip().lower() in allowed + from application.utils import roles + + return roles.has_role(email, roles.ADMIN) def verify_admin_allowlisted_api_token(raw: str) -> Optional[ApiKeyIdentity]: - """Path-agnostic Admin JWT check: signature + DB + LOGIN_ALLOWED_EMAILS.""" + """Path-agnostic admin check: signature + DB + admin role.""" identity = verify_api_token(raw) if identity is None: return None @@ -109,6 +105,16 @@ def verify_admin_allowlisted_api_token(raw: str) -> Optional[ApiKeyIdentity]: return identity +def verify_token_with_role(raw: str, needed: str) -> Optional[ApiKeyIdentity]: + """Signature + DB + the account currently has at least role ``needed`` (None otherwise).""" + from application.utils import roles + + identity = verify_api_token(raw) + if identity is None or not roles.has_role(identity.email, needed): + return None + return identity + + def bind_api_key_identity(identity: ApiKeyIdentity) -> None: g.api_key_jti = identity.jti g.api_key_user_id = identity.user_id diff --git a/application/utils/comprehension/answer.py b/application/utils/comprehension/answer.py index b615548a5..fc416964f 100644 --- a/application/utils/comprehension/answer.py +++ b/application/utils/comprehension/answer.py @@ -172,6 +172,61 @@ def _split_citations(content: str) -> Tuple[str, List[str]]: ] +TEMPORARY_PROBLEM_MESSAGE = ( + "I hit a temporary problem while writing this answer (the model service did not respond). " + "Please try again in a moment." +) + + +class AnswerUnavailable(Exception): + """The model call behind an answer failed; not the same as "no source answers this".""" + + +_MISSING_SYSTEM = ( + "You check whether an answer addresses everything a question asks. List what the question asks for, " + "then say which asked-for items the answer does NOT address at all. Be strict only about whole items " + 'that are absent; ignore style and detail. Reply with JSON only: {"missing": ["short phrase", ...]} ' + "(empty list when the answer covers everything)." +) + + +def missing_parts(question: str, answer: str, helper: Any) -> List[str]: + """Items the question asks for that the answer leaves out (cheap checker).""" + from application.utils.comprehension import extras + + try: + data = extras.parse_json( + helper(_MISSING_SYSTEM, f"Question: {question}\n\nAnswer: {answer}") + ) + return [str(m).strip() for m in (data.get("missing") or []) if str(m).strip()][ + :4 + ] + except Exception: # noqa: BLE001 - checker failure means "no second pass" + return [] + + +def _complete_pass( + question: str, + out: Tuple[str, List[str]], + cands: Sequence[Any], + complete: CompletionFn, +) -> Tuple[str, List[str]]: + from application.utils.comprehension import flags + + helper = flags.helper_llm() + if helper is None or not out or not out[0]: + return out + missing = missing_parts(question, out[0], helper) + if not missing: + return out + again = _answer( + f"{question}\n(The answer must also cover: {'; '.join(missing)})", + cands, + complete, + ) + return again if again and again[0] else out + + def answer_from_candidates( question: str, cands: Sequence[Any], complete: Optional[CompletionFn] ) -> Optional[Tuple[str, List[str]]]: @@ -188,9 +243,21 @@ def answer_from_candidates( got = _ANSWERS.get(ckey) if got and time.time() - got[0] < 3600: return got[1] - out = _answer_decomposed(question, complete) if flags.on("DECOMPOSE") else None - if out is None: - out = _answer(question, cands, complete) + try: + out = _answer_decomposed(question, complete) if flags.on("DECOMPOSE") else None + if out is None: + out = _answer(question, cands, complete) + if out is not None and flags.on("COMPLETE2"): + out = _complete_pass(question, out, cands, complete) + except ( + Exception + ) as exc: # noqa: BLE001 - provider 503/429/timeouts must not look like "not found" + import logging + + logging.getLogger(__name__).warning( + "comprehension answer failed: %s: %s", type(exc).__name__, str(exc)[:200] + ) + raise AnswerUnavailable(str(exc)[:200]) from exc if flags.on("CACHE") and out is not None: _ANSWERS[ckey] = (time.time(), out) if len(_ANSWERS) > 512: @@ -252,12 +319,25 @@ def _answer_from_whole_project( from application.utils.comprehension.retrieve import named_slugs from application.utils.comprehension.synthesize import PERSTD_RULE, SYSTEM - system = SYSTEM + from application.utils.comprehension import extras + from application.utils.comprehension.synthesize import ( + COUNT_SYSTEM, + NEG_RE, + NEG_RULE, + ) + + system = ( + COUNT_SYSTEM + if flags.on("COUNT") and extras.COUNT_RE.search(question) + else SYSTEM + ) + if flags.on("NEG") and NEG_RE.search(question): + system = system + NEG_RULE sources: List[tuple] = [] if flags.on("PERSTD") and len(named_slugs(question)) >= 2: sources = project_sources(question, max_projects=3) if sources: - system = SYSTEM + PERSTD_RULE + system = system + PERSTD_RULE if not sources: sources = project_sources(question) if not sources: @@ -269,7 +349,14 @@ def llm(system: str, user: str) -> str: ) answer, cites = synthesize(question, sources, llm, system=system) - return (answer, cites) if answer else None + if not answer: + return None + from application.utils.comprehension.retrieve import last_pick + + assumed = last_pick() + if assumed: + answer = f"Assuming you mean {assumed}: {answer}" + return answer, cites def think_effort() -> str: @@ -321,6 +408,10 @@ def _answer( temperature=0.0, # Thinking tokens bill as output and add seconds: off, or a small budget. extra_try_kwargs={"reasoning_effort": effort} if effort else None, + interactive=True, + fallback_model=os.environ.get( + "OWASP_ANSWER_FALLBACK_MODEL", "gemini/gemini-2.5-flash-lite" + ), ) complete = lambda msgs: call( msgs[0]["content"], msgs[1]["content"] @@ -385,9 +476,19 @@ def llm(system: str, user: str) -> str: if plan.get("parts"): ask = f"{ask}\n(Answer each part: " + " | ".join(plan["parts"]) + ")" count_mode = flags.on("COUNT") and bool(extras.COUNT_RE.search(ask)) - answer, cites = synthesize( - ask, sources, llm, system=COUNT_SYSTEM if count_mode else SYSTEM - ) + system = COUNT_SYSTEM if count_mode else SYSTEM + if flags.on("NEG"): + from application.utils.comprehension.synthesize import NEG_RE, NEG_RULE + + if NEG_RE.search(question): + system = system + NEG_RULE + if flags.on("PERSTD"): + from application.utils.comprehension.retrieve import named_slugs + from application.utils.comprehension.synthesize import PERSTD_RULE + + if len(named_slugs(question)) >= 2: + system = system + PERSTD_RULE + answer, cites = synthesize(ask, sources, llm, system=system) if answer and flags.on("VERIFY"): hl = flags.helper_llm() if hl and not extras.cheap_verify(answer, [t for t, _ in sources]): @@ -396,4 +497,13 @@ def llm(system: str, user: str) -> str: _, answer = extras.verify_answer(ask, answer, numbered, hl) except Exception: pass - return (answer, cites) if answer else None + if not answer: + return None + from application.utils.comprehension import authority + + for ( + url + ) in authority.featured_urls(): # admin-featured projects: show the closest match + if url not in cites: + cites = list(cites) + [url] + return answer, cites diff --git a/application/utils/comprehension/authority.py b/application/utils/comprehension/authority.py new file mode 100644 index 000000000..64678e53f --- /dev/null +++ b/application/utils/comprehension/authority.py @@ -0,0 +1,142 @@ +"""Admin-editable authority weights: favour or soften projects in answers. + +Stored as one JSON setting ``{repo slug: weight}`` in ``admin_runtime_setting``. Unlisted +projects are 1.0 (neutral). + +* Below 1: the project's chunk scores are multiplied by the weight before ranking (0 hides it). +* Above 1 ("featured"): ranking is left alone, so answers do not get worse, but the project's best + matching chunk is listed among the answer's references whenever it is reasonably relevant to the + question. The higher the weight, the less relevant a match may be. Never shown for abstentions. +""" + +from __future__ import annotations + +import json +import time +from contextvars import ContextVar +from typing import Any, Dict, List, Optional + +KEY = "COMPREHENSION_AUTHORITY" +MIN_WEIGHT, MAX_WEIGHT = 0.0, 3.0 +_CACHE: Dict[str, Any] = {"at": 0.0, "weights": {}} +_TTL = 30.0 + + +def validate(raw: Any) -> Dict[str, float]: + """Clean ``{slug: weight}``; neutral (1.0) entries are dropped. Raises ValueError.""" + if not isinstance(raw, dict): + raise ValueError("weights must be an object of slug -> number") + out: Dict[str, float] = {} + for slug, val in raw.items(): + if not isinstance(slug, str) or not slug.strip(): + raise ValueError("project slug must be a non-empty string") + if isinstance(val, bool) or not isinstance(val, (int, float, str)): + raise ValueError(f"weight for {slug} must be a number") + try: + w = float(val) + except ValueError: + raise ValueError(f"weight for {slug} must be a number") from None + if not MIN_WEIGHT <= w <= MAX_WEIGHT: + raise ValueError( + f"weight for {slug} must be between {MIN_WEIGHT} and {MAX_WEIGHT}" + ) + if abs(w - 1.0) > 1e-9: + out[slug.strip()] = round(w, 3) + return out + + +def weights(fresh: bool = False) -> Dict[str, float]: + """Current weights (30 s cache); {} when none are set or the database is unavailable.""" + if not fresh and time.time() - _CACHE["at"] < _TTL: + return _CACHE["weights"] + loaded: Dict[str, float] = {} + try: + from application.utils.admin_panel import runtime_settings + + raw = runtime_settings.get_setting(KEY) + if raw: + loaded = validate(json.loads(raw)) + except ( + Exception + ): # noqa: BLE001 - authority is a bias, never a reason to fail an answer + loaded = {} + _CACHE.update(at=time.time(), weights=loaded) + return loaded + + +def save(new: Dict[str, float], actor: Optional[str] = None) -> Dict[str, float]: + from application.utils.admin_panel import runtime_settings + + clean = validate(new) + runtime_settings.set_setting(KEY, json.dumps(clean, sort_keys=True), actor=actor) + _CACHE.update(at=0.0, weights={}) + return clean + + +_FEATURED: "ContextVar[List[str]]" = ContextVar("authority_featured", default=[]) + + +def reset_featured() -> None: + _FEATURED.set([]) + + +def featured_urls() -> List[str]: + return list(_FEATURED.get()) + + +def tolerance(weight: float) -> float: + """How far below the best score a featured project's chunk may fall and still be listed.""" + return max(0.02, min(0.2, 0.08 * (weight - 1.0) / 0.15)) + + +def set_featured(urls: List[str]) -> None: + _FEATURED.set(list(urls)) + + +def bias(hits: List[Any], slug_of) -> List[Any]: + """Multiply each hit's score by its project's weight, drop weight-0 projects, re-sort.""" + w = weights() + if not w or all(v >= 1.0 for v in w.values()): + return hits + out = [] + for h in hits: + factor = min(1.0, w.get(slug_of(h) or "", 1.0)) + if factor <= 0: + continue + h.score = float(h.score) * factor + out.append(h) + out.sort(key=lambda h: -float(h.score)) + return out + + +def diversify(hits: List[Any], slug_of, projects: int = 3, per: int = 5) -> List[Any]: + """Make sure the best ``projects`` projects each get up to ``per`` hits at the front.""" + chosen: List[str] = [] + for h in hits: + s = slug_of(h) + if s and s not in chosen: + chosen.append(s) + if len(chosen) >= projects: + break + if len(chosen) < 2: + return hits + taken: Dict[str, int] = {} + front: List[Any] = [] + for h in hits: + s = slug_of(h) + if s in chosen and taken.get(s, 0) < per: + taken[s] = taken.get(s, 0) + 1 + front.append(h) + seen = {id(h) for h in front} + return front + [h for h in hits if id(h) not in seen] + + +def project_table(slugs: List[str]) -> List[Dict[str, Any]]: + from application.utils.comprehension.retrieve import alias_phrase + + w = weights(fresh=True) + rows = [ + {"slug": s, "name": alias_phrase(s), "weight": w.get(s, 1.0)} + for s in sorted(set(slugs) | set(w)) + ] + return rows diff --git a/application/utils/comprehension/embedcache.py b/application/utils/comprehension/embedcache.py index 179e2e135..46ec647f4 100644 --- a/application/utils/comprehension/embedcache.py +++ b/application/utils/comprehension/embedcache.py @@ -23,8 +23,20 @@ def _store(text: str, vec: Any) -> Any: return vec +def _quick(fn: Callable[[str], Any]) -> Callable[[str], Any]: + """Embedding calls a person waits on retry transient provider errors quickly.""" + + def _call(text: str) -> Any: + from application.prompt_client import litellm_router + + return litellm_router.quick_retry(lambda: fn(text), context="query embedding") + + return _call + + def cached(fn: Callable[[str], Any]) -> Callable[[str], Any]: """Wrap an embedding function with the shared cache (also joins in-flight prefetches).""" + fn = _quick(fn) def _embed(text: str) -> Any: with _LOCK: @@ -47,7 +59,8 @@ def prefetch(text: str, fn: Callable[[str], Any]) -> None: with _LOCK: if text in _CACHE or text in _PENDING: return - fut = _POOL.submit(lambda: _store(text, fn(text))) + quick = _quick(fn) + fut = _POOL.submit(lambda: _store(text, quick(text))) _PENDING[text] = fut def _done(_f: Any, t: str = text) -> None: diff --git a/application/utils/comprehension/flags.py b/application/utils/comprehension/flags.py index 1cb2f98aa..c39ab90f8 100644 --- a/application/utils/comprehension/flags.py +++ b/application/utils/comprehension/flags.py @@ -24,8 +24,14 @@ "PROJROUTE", # no project named: pick one by similarity to project overviews "COMPLETE", # answer every part and every named standard; include all stated items "PERSTD", # questions naming 2-3 projects: read each whole, answer each standard separately + "PICK", # no project named: a cheap model picks it from the best-matching projects (or says unsure) + "COMPLETE2", # second pass: a cheap checker finds unanswered parts, the reader answers them + "NEG", # negation questions (NOT / except): test every option against the sources before choosing + "PREFIXCACHE", # sources first, question last, so the provider can reuse the cached project text + "WIDEMATCH", # whole-project reading can match any repo with raw chunks, not only digested ones "EDITION", # answer for the edition/year/version the question names; never mix editions "ROUTE", # people/count questions: whole-project read with a stronger reader model + "TOPPROJECTS", # no project named: the 3 best-matching projects each contribute sources "DECOMPOSE", # multi-part questions: answer each part separately, then combine ) diff --git a/application/utils/comprehension/retrieve.py b/application/utils/comprehension/retrieve.py index 9bdb69044..66f4020db 100644 --- a/application/utils/comprehension/retrieve.py +++ b/application/utils/comprehension/retrieve.py @@ -8,6 +8,7 @@ import os import re +from contextvars import ContextVar import time from typing import Any, List @@ -48,16 +49,33 @@ def matched_slugs(question: str, slugs: List[str], limit: int = 3) -> List[str]: def _known_slugs(conn) -> List[str]: from sqlalchemy import text as sql_text - if time.time() - _REPO_CACHE["at"] < 600 and _REPO_CACHE["slugs"]: + from application.utils.comprehension import flags + + wide = flags.on("WIDEMATCH") + cache_key = "wide" if wide else "digested" + if ( + _REPO_CACHE.get("key") == cache_key + and time.time() - _REPO_CACHE["at"] < 600 + and _REPO_CACHE["slugs"] + ): return _REPO_CACHE["slugs"] - rows = conn.execute( - sql_text( - "SELECT DISTINCT embeddings_url FROM embeddings " - "WHERE embeddings_url LIKE 'owasp-agent://evidence/comprehension-%-overview'" - ) - ).fetchall() - slugs = sorted({m.group(1) for (u,) in rows if (m := _SLUG_RE.match(u))}) - _REPO_CACHE.update(at=time.time(), slugs=slugs) + if wide: + rows = conn.execute( + sql_text( + "SELECT DISTINCT substring(embeddings_url from 'comprehension-(.+?)-raw-') " + "FROM embeddings WHERE embeddings_url LIKE 'owasp-agent://evidence/comprehension-%-raw-%'" + ) + ).fetchall() + slugs = sorted({r[0] for r in rows if r[0]}) + else: + rows = conn.execute( + sql_text( + "SELECT DISTINCT embeddings_url FROM embeddings " + "WHERE embeddings_url LIKE 'owasp-agent://evidence/comprehension-%-overview'" + ) + ).fetchall() + slugs = sorted({m.group(1) for (u,) in rows if (m := _SLUG_RE.match(u))}) + _REPO_CACHE.update(at=time.time(), slugs=slugs, key=cache_key) return slugs @@ -119,6 +137,15 @@ def _dense( from application.utils.owasp_agent.lane_similarity import cosine_exprs dist, score = cosine_exprs() + if limit > 40: # the HNSW index returns at most ef_search (default 40) rows + try: + conn.execute( + sql_text(f"SET LOCAL hnsw.ef_search = {min(1000, int(limit))}") + ) + except ( + Exception + ): # noqa: BLE001 - non-pgvector / older pgvector: keep the default + pass rows = conn.execute( sql_text(_DENSE_SQL.format(where=where, dist=dist, score=score)), {"q": q, "channel": CHANNEL_AGENT, "limit": limit, **(params or {})}, @@ -176,6 +203,51 @@ def _expand(conn: Any, hits: List[LaneHit]) -> None: h.content = merge_neighbours(h.content, prev, nxt) +def _feature_projects(conn: Any, q: str, dense: List[LaneHit]) -> None: + """Remember each featured project's best chunk when it is close enough to the best match.""" + from application.utils.comprehension import authority + + try: + featured = {s: w for s, w in authority.weights().items() if w > 1.0} + except Exception: # noqa: BLE001 + return + if not featured or not dense: + return + best = max(float(h.score) for h in dense) + urls = authority.featured_urls() + for slug, w in featured.items(): + top = _dense( + conn, + q, + "embeddings_url LIKE :pat", + 1, + {"pat": f"owasp-agent://evidence/comprehension-{slug}-%"}, + ) + if top and float(top[0].score) >= best - authority.tolerance(w): + _, _, tail = top[0].content.partition("\n\nCitations:") + cite = next( + (c.strip() for c in tail.split(",") if c.strip().startswith("http")), "" + ) + if cite and cite not in urls: + urls.append(cite) + authority.set_featured(urls) + + +def _slug_of(hit: LaneHit) -> str: + m = _SLUG_RE.match(hit.url or "") + return m.group(1) if m else "" + + +def _apply_authority(hits: List[LaneHit]) -> List[LaneHit]: + """Admin authority weights, then (TOPPROJECTS) the best 3 projects each get a share.""" + from application.utils.comprehension import authority, flags + + hits = authority.bias(hits, _slug_of) + if flags.on("TOPPROJECTS"): + hits = authority.diversify(hits, _slug_of, projects=3, per=5) + return hits + + def fetch_comprehension_hits( question: str, queries: List[str] | None = None ) -> List[LaneHit]: @@ -198,6 +270,9 @@ def fetch_comprehension_hits( from application.utils.comprehension.shared import shared_handler handler = shared_handler(database) + from application.utils.comprehension import authority as _authority + + _authority.reset_featured() texts = [question] + [q for q in (queries or []) if q and q != question] k = retrieval_k() wide = k * 2 if flags.on("RERANK") else k @@ -210,8 +285,15 @@ def fetch_comprehension_hits( vecs = embed_many(texts, handler.get_text_embeddings) for t, vec in zip(texts, vecs): q = to_pgvector_literal(list(vec)) - dense = _dense(conn, q, docs_where(), wide) + dense = _dense( + conn, + q, + docs_where(), + max(wide, 300) if flags.on("TOPPROJECTS") else wide, + ) rankings.append(dense) + if t == question: + _feature_projects(conn, q, dense) if flags.on("HYBRID"): keyword_hits += _keyword(conn, t, wide) scoped = matched_slugs(t, slugs) @@ -246,6 +328,7 @@ def fetch_comprehension_hits( hits = rrf(rankings, key=lambda h: h.url) if keyword_hits: hits = supplement_with_keyword(hits, keyword_hits) + hits = _apply_authority(hits) if flags.on("EXPAND"): _expand(conn, hits[: max(k, 12)]) return hits @@ -364,6 +447,75 @@ def route_project(conn: Any, question: str) -> List[str]: return [m.group(1)] if m else [] +PICK_SYSTEM = ( + "You decide which single OWASP project a question is about. Candidates are numbered with a " + 'snippet of their text. Answer with JSON only: {"project": , ' + '"confidence": <0 to 1>}. Use null when the question could be about several candidates, ' + "about none of them, or does not make clear which project it means." +) +_PICK_STATE: "ContextVar[str]" = ContextVar("comprehension_pick", default="") + + +def last_pick() -> str: + """Readable name of the project chosen by PICK for the current question ('' if none).""" + return _PICK_STATE.get() + + +def pick_candidates(conn: Any, question: str, limit: int = 8) -> List[tuple]: + """(slug, snippet) of the best-matching distinct projects, best first.""" + from application.database import db as db_mod + from application.database.pgvector_utils import to_pgvector_literal + from application.utils.comprehension.embedcache import embed_many + from application.utils.comprehension.shared import shared_handler + + handler = shared_handler(db_mod.Node_collection()) + vec = embed_many([question], handler.get_text_embeddings)[0] + hits = _dense(conn, to_pgvector_literal(list(vec)), docs_where(), 60) + seen: dict = {} + for h in hits: + m = _SLUG_RE.match(h.url or "") + if not m or m.group(1) in seen: + continue + body = ( + (h.content or "") + .split("\n\nCitations:")[0] + .split("\n\nRelated questions:")[0] + ) + seen[m.group(1)] = " ".join(body.split())[:220] + if len(seen) >= limit: + break + return list(seen.items()) + + +def pick_project(conn: Any, question: str, llm: Any = None) -> str: + """Slug of the project the question is about, chosen by a cheap model; '' when unsure.""" + from application.utils.comprehension import extras, flags + + llm = llm or flags.helper_llm() + if llm is None: + return "" + cands = pick_candidates(conn, question) + if len(cands) < 2: + return "" + menu = "\n".join( + f"{i}. {alias_phrase(slug)}: {snippet}" + for i, (slug, snippet) in enumerate(cands, 1) + ) + try: + data = extras.parse_json( + llm(PICK_SYSTEM, f"Question: {question}\n\nCandidates:\n{menu}") + ) + n = int(data.get("project")) + conf = float(data.get("confidence") or 0) + except Exception: # noqa: BLE001 - null / malformed means "unsure" + return "" + if conf < float(os.environ.get("OWASP_CX_PICK_MIN", "0.7")) or not 1 <= n <= len( + cands + ): + return "" + return cands[n - 1][0] + + def named_slugs(question: str, limit: int = 4) -> List[str]: """Repo slugs of the projects the question names (up to ``limit``); [] on any failure.""" try: @@ -385,11 +537,17 @@ def project_sources(question: str, max_projects: int | None = None) -> List[tupl limit = max_projects or longctx_max_projects() with sqla.session.get_bind().connect() as conn: slugs = matched_slugs(question, _known_slugs(conn), limit=limit + 1) + _PICK_STATE.set("") if not slugs: from application.utils.comprehension import flags if flags.on("PROJROUTE"): slugs = route_project(conn, question) + elif flags.on("PICK"): + chosen = pick_project(conn, question) + if chosen: + slugs = [chosen] + _PICK_STATE.set(alias_phrase(chosen).title()) if not slugs or len(slugs) > limit: return [] rows: List[tuple] = [] diff --git a/application/utils/comprehension/shared.py b/application/utils/comprehension/shared.py index 6ecb0ec5a..f62bb1977 100644 --- a/application/utils/comprehension/shared.py +++ b/application/utils/comprehension/shared.py @@ -10,6 +10,14 @@ _HANDLER: Any = None +def _interactive(handler: Any) -> Any: + """This handler serves people waiting for answers: the interactive layer + (litellm_router.quick_retry / resilient_completion) does the retrying, so the bulk-job sleeps + inside the handler are switched off.""" + handler._llm_max_retries = 0 + return handler + + def shared_handler(database: Any) -> Any: """A lazily created, cached PromptHandler(load_all_embeddings=False).""" from application.prompt_client import prompt_client @@ -21,10 +29,14 @@ def shared_handler(database: Any) -> Any: "no", "off", ): - return prompt_client.PromptHandler(database, load_all_embeddings=False) + return _interactive( + prompt_client.PromptHandler(database, load_all_embeddings=False) + ) with _LOCK: if _HANDLER is None: - _HANDLER = prompt_client.PromptHandler(database, load_all_embeddings=False) + _HANDLER = _interactive( + prompt_client.PromptHandler(database, load_all_embeddings=False) + ) return _HANDLER @@ -57,3 +69,25 @@ def warm_up() -> bool: return True except Exception: # noqa: BLE001 - warm-up must never stop the server return False + + +def start_background_warm_up(app: Any) -> bool: + """Warm a freshly started web worker in the background so the first chat is not the slow one. + + Opt-in with OWASP_WARMUP_WEB=1 (command-line jobs that also call create_app must not start it). + """ + if (os.environ.get("OWASP_WARMUP_WEB", "") or "").strip().lower() not in ( + "1", + "true", + "yes", + "on", + ): + return False + import threading + + def _run() -> None: + with app.app_context(): + warm_up() + + threading.Thread(target=_run, name="opencre-web-warmup", daemon=True).start() + return True diff --git a/application/utils/comprehension/sources.py b/application/utils/comprehension/sources.py index 5c8b2e962..501a389e4 100644 --- a/application/utils/comprehension/sources.py +++ b/application/utils/comprehension/sources.py @@ -184,3 +184,59 @@ def supplement_chunks( if len(piece) >= 60: out.append((rel, i, piece[:6000])) return out + + +_WIKI_EXT = ( + ".md", + ".markdown", + ".txt", + ".textile", + ".rdoc", + ".mediawiki", + ".asciidoc", + ".rst", + ".org", +) + + +def wiki_chunks(repo: str, clone_dir: str) -> List[Tuple[str, int, str]]: + """Chunks of the repository's GitHub wiki (a separate git repo); [] when it has none.""" + import os + import subprocess + from pathlib import Path + + dest = Path(clone_dir) / (repo.replace("/", "__") + "__wiki") + if not dest.exists(): + done = subprocess.run( + [ + "git", + "clone", + "--depth", + "1", + "-q", + f"https://github.com/{repo}.wiki.git", + str(dest), + ], + capture_output=True, + timeout=120, + env={**os.environ, "GIT_TERMINAL_PROMPT": "0"}, + ) + if done.returncode != 0: + return [] + out: List[Tuple[str, int, str]] = [] + for f in sorted(dest.rglob("*")): + rel = str(f.relative_to(dest)) + if ( + not f.is_file() + or ".git/" in rel + or rel.startswith(".git") + or not rel.lower().endswith(_WIKI_EXT) + ): + continue + text = f.read_text(errors="ignore") + if len(text.strip()) < 80: + continue + for i, piece in enumerate(split_markdown(text)): + if len(piece) >= 60: + out.append((f"wiki/{rel}", i, piece[:6000])) + return out diff --git a/application/utils/comprehension/synthesize.py b/application/utils/comprehension/synthesize.py index 08c522d37..36275a0ca 100644 --- a/application/utils/comprehension/synthesize.py +++ b/application/utils/comprehension/synthesize.py @@ -4,7 +4,7 @@ import re from contextvars import ContextVar -from typing import Callable, List, Sequence, Tuple +from typing import Callable, Dict, List, Optional, Sequence, Tuple LlmFn = Callable[[str, str], str] @@ -13,6 +13,45 @@ _EVIDENCE: "ContextVar[list]" = ContextVar("comprehension_evidence", default=[]) +# The chat's earlier turns for the question being answered. The answer call sends them as real +# user/assistant messages (native multi-turn), so the model reads them as background, not as questions. +_HISTORY: "ContextVar[list]" = ContextVar("comprehension_history", default=[]) +HISTORY_RULE = ( + " Earlier messages in this conversation are background only: answer ONLY the final question, " + "do not answer the earlier questions again." +) +MAX_HISTORY_TURNS = 6 +MAX_HISTORY_CHARS = 600 + + +def set_history(history) -> None: + """Prepare the earlier chat turns (last few, trimmed, alternating roles) for the answer call.""" + from application.utils.chat_history import normalize_chat_messages + + turns = [] + for m in normalize_chat_messages(history)[-MAX_HISTORY_TURNS:]: + content = m["content"][:MAX_HISTORY_CHARS] + if turns and turns[-1]["role"] == m["role"]: + turns[-1]["content"] += "\n" + content # providers want alternating roles + else: + turns.append({"role": m["role"], "content": content}) + while turns and turns[0]["role"] != "user": + turns.pop(0) + _HISTORY.set(turns) + + +def get_history() -> list: + return list(_HISTORY.get()) + + +def with_history(messages: list) -> list: + """[system, ...earlier turns, final user message]: native multi-turn for the answer call.""" + turns = get_history() + if not turns or len(messages) < 2: + return list(messages) + return [messages[0], *turns, *messages[1:]] + + def reset_evidence() -> None: _EVIDENCE.set([]) @@ -21,8 +60,61 @@ def take_evidence() -> list: return list(_EVIDENCE.get()) +_MD_LINK = re.compile(r"\[([^\[\]]+)\]\(\s*https?://[^)]*\)") +_MARK_LINK = re.compile(r"\[\[(\d+)\]\]\(\s*https?://[^)]*\)") +_DOUBLE_MARK = re.compile(r"\[\[(\d+)\]\]") +_BARE_URL = re.compile(r"https?://[^\s)\]>]+") + + +def clean_answer_text(text: str) -> str: + """Keep the model's answer free of links: only ``[n]`` markers may point at sources. + + Sources contain markdown (``[text](url)``) and the model sometimes copies it, producing + ``[[7]](https://...)`` that the UI cannot link and a raw URL that should stay hidden. + The references list under the answer carries the real links. + """ + out = _MARK_LINK.sub(lambda m: f"[{m.group(1)}]", text) + out = _MD_LINK.sub( + lambda m: f"[{m.group(1)}]" if m.group(1).isdigit() else m.group(1), out + ) + out = _DOUBLE_MARK.sub(lambda m: f"[{m.group(1)}]", out) + out = _BARE_URL.sub("", out) + return re.sub(r"[ \t]+([.,;:])", r"\1", out) + + +def renumber_citations( + answer: str, + hits: Sequence[Tuple[str, List[str]]], + used: Sequence[int], + cites: Sequence[str], +) -> Tuple[str, Dict[int, int]]: + """Make the [n] markers match the returned citation list (so a UI can link [n] to cites[n-1]). + + The model cites the numbered sources it was shown ([3], [17], ...), but callers receive only the URLs + of the sources it used. Each marker becomes the 1-based position of its source's first URL in ``cites``; + several chunks of one file share a number, repeated markers collapse ([2][2] -> [2]) and markers of + sources without a URL are dropped. + """ + number: Dict[int, int] = {} + for i in used: + urls = hits[i - 1][1] + if urls and urls[0] in cites: + number[i] = list(cites).index(urls[0]) + 1 + + def swap(m: "re.Match[str]") -> str: + n = number.get(int(m.group(1))) + return f"[{n}]" if n else "" + + out = re.sub(r"\[(\d+)\]", swap, answer) + out = re.sub(r"(\[\d+\])(?:\s*\1)+", r"\1", out) # [2][2] -> [2] + out = re.sub(r"[ \t]+([.,;:!?])", r"\1", out) + return out, number + + def _record_evidence( - hits: Sequence[Tuple[str, List[str]]], used: Sequence[int] + hits: Sequence[Tuple[str, List[str]]], + used: Sequence[int], + number: Optional[Dict[int, int]] = None, ) -> None: from application.utils.comprehension.toolkit import body_of, path_of @@ -33,7 +125,7 @@ def _record_evidence( snippet = body.split(": ", 1)[-1] if path_of(body) else body items.append( { - "n": i, + "n": number.get(i) if number is not None else i, "path": path_of(body), "snippet": snippet.strip()[:500], "citations": list(cites)[:3], @@ -85,6 +177,16 @@ def _record_evidence( ) +NEG_RE = re.compile( + r"\bnot\b|n't\b|\bexcept\b|\bexcluding\b|\bnone of\b|\bnever\b|\bwithout\b", re.I +) +NEG_RULE = ( + " The question asks what is NOT / EXCEPT / excluded. First list every candidate option the question " + "or the sources give, then test each one against the sources (stated: yes/no), and only then name the " + "option(s) the sources do not support. Do not answer from the first sentence you find." +) + + def synthesize( question: str, hits: Sequence[Tuple[str, List[str]]], @@ -95,6 +197,8 @@ def synthesize( """``hits`` = (text, citations). Returns (answer, citations used); ("", []) if unknown.""" from application.utils.comprehension import flags + if get_history(): + system += HISTORY_RULE if flags.on("EDITION"): system += EDITION_RULE if flags.on("COMPLETE"): @@ -103,16 +207,22 @@ def synthesize( f"[{i + 1}] {t.split(chr(10) + chr(10) + 'Citations:')[0]}" for i, (t, _) in enumerate(hits) ) - out = (llm(system, f"Question: {question}\n\nSources:\n{numbered}") or "").strip() + if flags.on("PREFIXCACHE"): + # Same project text first: providers cache the longest shared prefix across questions. + user = f"Sources:\n{numbered}\n\nQuestion: {question}" + else: + user = f"Question: {question}\n\nSources:\n{numbered}" + out = clean_answer_text((llm(system, user) or "").strip()) if not out or out.upper().startswith("UNKNOWN"): return "", [] used = sorted( {int(n) for n in re.findall(r"\[(\d+)\]", out) if 0 < int(n) <= len(hits)} ) - _record_evidence(hits, used) cites: List[str] = [] for i in used: for c in hits[i - 1][1]: if c not in cites: cites.append(c) + out, number = renumber_citations(out, hits, used, cites) + _record_evidence(hits, used, number) return out, cites diff --git a/application/utils/comprehension/talks_source.py b/application/utils/comprehension/talks_source.py deleted file mode 100644 index 93891178b..000000000 --- a/application/utils/comprehension/talks_source.py +++ /dev/null @@ -1,46 +0,0 @@ -"""Talk catalogue (title/speaker/date/length per video) as retrievable documents.""" - -from __future__ import annotations - -from typing import Dict, List, Tuple - - -def parse_catalogue(text: str) -> List[Dict[str, str]]: - """Lines 'id|YYYYMMDD|seconds|title'; blank or malformed lines are skipped.""" - out: List[Dict[str, str]] = [] - for line in (text or "").splitlines(): - parts = line.rstrip("\n").split("|", 3) - if len(parts) != 4 or not parts[0].strip() or not parts[3].strip(): - continue - vid, date, secs, title = (p.strip() for p in parts) - out.append({"id": vid, "date": date, "seconds": secs, "title": title}) - return out - - -def _fmt_date(d: str) -> str: - return f"{d[:4]}-{d[4:6]}-{d[6:8]}" if len(d) == 8 and d.isdigit() else "unknown" - - -def _fmt_len(s: str) -> str: - return f"{int(s) // 60} min" if s.isdigit() else "?" - - -def catalogue_chunks( - talks: List[Dict[str, str]], per_chunk: int = 30 -) -> List[Tuple[str, int, str]]: - """Newest first; each chunk lists talks as 'date (approx) | length | title | link'.""" - ordered = sorted(talks, key=lambda t: t["date"], reverse=True) - chunks: List[Tuple[str, int, str]] = [] - for i, start in enumerate(range(0, len(ordered), per_chunk)): - part = ordered[start : start + per_chunk] - lines = [ - f"{_fmt_date(t['date'])} (approximate) | {_fmt_len(t['seconds'])} | " - f"{t['title']} | https://youtu.be/{t['id']}" - for t in part - ] - header = ( - "OWASP Foundation YouTube channel talk catalogue (public talks; speaker names are in the title; " - "dates are approximate, taken from the channel listing):\n" - ) - chunks.append(("talks/catalogue.md", i, header + "\n".join(lines))) - return chunks diff --git a/application/utils/owasp_agent/candidate_picker.py b/application/utils/owasp_agent/candidate_picker.py index 2b55cd748..61e483c1d 100644 --- a/application/utils/owasp_agent/candidate_picker.py +++ b/application/utils/owasp_agent/candidate_picker.py @@ -212,18 +212,21 @@ def default_completion_fn() -> Optional[CompletionFn]: database = db_mod.Node_collection() from application.utils.comprehension.shared import shared_handler + from application.utils.comprehension.synthesize import with_history handler = shared_handler(database) def _complete(messages: Sequence[dict]) -> str: - resp = litellm_router.completion( + # A person is waiting: quick retries, a cheaper fallback model and a circuit breaker. + resp = litellm_router.resilient_completion( model=handler.chat_model, - messages=list(messages), + fallback_model=os.environ.get( + "OWASP_ANSWER_FALLBACK_MODEL", "gemini/gemini-2.5-flash-lite" + ), + messages=with_history(list(messages)), client=handler._litellm, - context="OWASP agent candidate pick", + context="OWASP agent answer", temperature=0, - max_retries=handler._llm_max_retries, - retry_sleep_seconds=handler._llm_retry_sleep_seconds, ) return litellm_router.extract_content_text(resp) diff --git a/application/utils/owasp_agent/router.py b/application/utils/owasp_agent/router.py index fcaf38759..3a99c0e23 100644 --- a/application/utils/owasp_agent/router.py +++ b/application/utils/owasp_agent/router.py @@ -275,6 +275,12 @@ def handle( """ if not is_owasp_agent_enabled(): return None + if not _clause: + from application.utils.comprehension.synthesize import set_history + + set_history( + history + ) # native multi-turn for the answer call (background, not a question) text = fix_key_typos( (prompt or "").strip()[:MAX_PROMPT_CHARS].replace("\u2019", "'") ) @@ -372,7 +378,25 @@ def handle( from application.utils.comprehension.synthesize import take_evidence comp_hits = [c for c in cands if is_comprehension_hit(c)] - comp = answer_from_candidates(normalized, cands, default_completion_fn()) + from application.utils.comprehension.answer import ( + TEMPORARY_PROBLEM_MESSAGE, + AnswerUnavailable, + ) + + try: + comp = answer_from_candidates( + normalized, cands, default_completion_fn() + ) + except AnswerUnavailable: + return format_chat_response( + QueryResult( + ok=False, + kind="error", + message=TEMPORARY_PROBLEM_MESSAGE, + citations=[], + data={"pick_reason": "comprehension_unavailable"}, + ) + ) if comp is not None: return format_chat_response( QueryResult( @@ -2310,10 +2334,103 @@ def _clean_project_name(raw: str) -> str: return name -def _effective_prompt( - prompt: str, history: Optional[List[Dict[str, Any]]] = None +_FOLLOW_UP_START = re.compile( + r"^(and|but|also|so|then|what about|how about|why|how so|what else|more|go on|tell me more|" + r"can you|could you|please|ok(ay)?|yes|no|sure)\b", + re.IGNORECASE, +) +_PRONOUN = re.compile( + r"\b(it|its|they|them|their|this|that|these|those|he|she|there|one)\b", + re.IGNORECASE, +) +_STANDALONE_SYSTEM = ( + "Rewrite the user's LAST message as one standalone question, using the earlier conversation ONLY to " + "fill in what the last message refers to (the subject, project, standard or place it leaves out). " + "Do not answer anything. Do not add, repeat or combine the earlier questions. Keep it short. " + "If the last message already stands on its own, return it unchanged. Reply with the question only." +) + + +def needs_context(prompt: str) -> bool: + """True when the message cannot be understood without the conversation (short, elliptical, pronouns).""" + text = (prompt or "").strip() + words = text.split() + if not words: + return False + if len(words) <= 5 or _FOLLOW_UP_START.match(text): + return True + if _PRONOUN.search(text): + # A pronoun with a named subject elsewhere ("Does ASVS cover it?") still needs the referent, + # but a full sentence naming a project or standard is treated as standalone. + named = [ + w for w in words[1:] if w[:1].isupper() or any(ch.isdigit() for ch in w) + ] + return not named + return False + + +def _history_text(history: Optional[List[Dict[str, Any]]], limit: int = 4) -> str: + turns = [] + for turn in history or []: + if not isinstance(turn, dict): + continue + role = str(turn.get("role") or "").lower() + content = " ".join(str(turn.get("content") or "").split()) + if role in ("user", "assistant") and content: + turns.append(f"{role}: {content[:400]}") + return "\n".join(turns[-limit:]) + + +def standalone_question( + prompt: str, + history: Optional[List[Dict[str, Any]]] = None, + llm: Optional[Any] = None, ) -> str: - """Join recent user turns so clarify follow-ups keep prior slots/context.""" + """The question to answer: the current message, made self-contained when it depends on the chat. + + The conversation is background, never a second question. A message that already stands on its own is + returned untouched (history ignored). A follow-up ("What are some key risks?") is rewritten by a cheap + model into a standalone question; if that fails, the current message is used on its own. + ``OWASP_HISTORY_MODE=concat`` restores the old behaviour (earlier user turns glued in front). + """ + mode = (os.environ.get("OWASP_HISTORY_MODE") or "rewrite").strip().lower() + if mode == "concat": + return _concat_prompt(prompt, history) + text = (prompt or "").strip() + context = _history_text(history) + if not context or not needs_context(text): + return text + try: + if llm is None: + from application.prompt_client import litellm_router + + llm = litellm_router.system_user_fn( + os.environ.get("OWASP_REWRITE_MODEL", "gemini/gemini-2.5-flash-lite"), + temperature=0.0, + extra_try_kwargs={"reasoning_effort": "disable"}, + interactive=True, + ) + out = " ".join( + ( + llm( + _STANDALONE_SYSTEM, + f"Conversation:\n{context}\n\nLast message: {text}", + ) + or "" + ).split() + ) + except ( + Exception + ): # noqa: BLE001 - on any failure answer the current message alone, never the history + return text + out = out.strip().strip('"') + if not out or len(out) > 400 or out.count("?") > 1: + return text # empty, rambling, or several questions glued together: not a standalone question + return out + + +def _concat_prompt(prompt: str, history: Optional[List[Dict[str, Any]]] = None) -> str: + """Old behaviour: join recent user turns so clarify follow-ups keep prior slots/context.""" prior: List[str] = [] for turn in history or []: if not isinstance(turn, dict): @@ -2328,3 +2445,9 @@ def _effective_prompt( if not prior: return prompt return (" ".join(prior[-3:]) + "\n" + prompt).strip() + + +def _effective_prompt( + prompt: str, history: Optional[List[Dict[str, Any]]] = None +) -> str: + return standalone_question(prompt, history) diff --git a/application/utils/owasp_agent/service.py b/application/utils/owasp_agent/service.py new file mode 100644 index 000000000..c268f24fd --- /dev/null +++ b/application/utils/owasp_agent/service.py @@ -0,0 +1,177 @@ +"""One entry point for "answer this OWASP question", shared by the web chatbot and MCP ``ask``. + +Both call :func:`route_question`, so routing, the answer pipeline, error handling and the +evidence returned are the same. The only caller-specific input is ``history`` (the web chat +keeps conversation memory; MCP callers may pass it too). +""" + +from __future__ import annotations + +import logging +import os +from typing import Any, Dict, List, Optional + +logger = logging.getLogger(__name__) + + +class OwaspAgentUnavailable(Exception): + """The OWASP agent failed on a question that must not fall through to CRE retrieval.""" + + +def _cre_attach_threshold() -> float: + try: + return float(os.environ.get("OWASP_CRE_ATTACH_THRESHOLD", "0.6")) + except ValueError: + return 0.6 + + +_NAME_MATCH_BONUS = 0.03 +_STOP = frozenset( + "how do does did can could should would what which who why when where is are was the " + "a an and or of to for in on with my our your i we it this that about from by be as at " + "prevent avoid stop protect against secure securely best way ways make using use".split() +) + + +def _significant_words(text: str) -> set: + """Lower-case content words (3+ letters) of a question or CRE name.""" + import re + + return { + w + for w in re.findall(r"[a-z0-9]+", (text or "").lower()) + if len(w) >= 3 and w not in _STOP + } + + +def closest_cre_rows( + prompt: str, limit: int = 3, margin: float = 0.03 +) -> List[Dict[str, Any]]: + """The CREs closest to the question: the best match plus any within ``margin`` of it. + + OpenCRE chat exists to connect a control question to its Common Requirement(s), so every + evidence answer carries them even when the answer text comes from other sources. Near-ties + (e.g. "Escape output against XSS" vs "XSS protection") are all shown instead of picking one. + ``[]`` when none is close enough or on any failure. + """ + try: + from sqlalchemy import text as sql_text + + from application import sqla + from application.database import db as db_mod + from application.database.pgvector_utils import to_pgvector_literal + from application.utils.comprehension.embedcache import embed_many + from application.utils.comprehension.shared import shared_handler + from application.utils.owasp_agent.lane_similarity import cosine_exprs + + database = db_mod.Node_collection() + handler = shared_handler(database) + vec = embed_many([prompt], handler.get_text_embeddings)[0] + dist, score = cosine_exprs() + with sqla.session.get_bind().connect() as conn: + found = conn.execute( + sql_text( + f"SELECT cre_id, {score} AS s FROM embeddings " + "WHERE doc_type = 'CRE' AND embedding_vec IS NOT NULL " + f"ORDER BY {dist} LIMIT :n" + ), + {"q": to_pgvector_literal(list(vec)), "n": limit * 3}, + ).fetchall() + if not found or float(found[0].s) < _cre_attach_threshold(): + return [] + words = _significant_words(prompt) + scored = [] + seen = set() + for r in found: + if r.cre_id in seen: + continue + seen.add(r.cre_id) + cre = database.get_cre_by_db_id(r.cre_id) + if not cre: + continue + bonus = _NAME_MATCH_BONUS if words & _significant_words(cre.name) else 0.0 + scored.append((float(r.s) + bonus, cre)) + scored.sort(key=lambda t: -t[0]) + best = scored[0][0] if scored else 0.0 + rows = [ + cre.shallow_copy().todict() for sc, cre in scored if sc >= best - margin + ][:limit] + return rows + except Exception: # noqa: BLE001 - an extra reference must never fail an answer + logger.info("closest CRE lookup failed", exc_info=True) + return [] + + +def attach_closest_cre(prompt: str, response: Dict[str, Any]) -> Dict[str, Any]: + """Add the closest CRE(s) to an evidence answer's references and citations (CRE-first bias).""" + if not isinstance(response, dict) or not response.get("accurate", True): + return response + meta = response.get("owasp_agent") + if not isinstance(meta, dict) or meta.get("kind") != "evidence": + return response + rows = closest_cre_rows(prompt) + if not rows: + return response + out = dict(response) + table = out.get("table") + out["table"] = rows + [r for r in (table if isinstance(table, list) else []) if r] + cites = list(meta.get("citations") or []) + for row in rows: + cre_id = str(row.get("id") or "") + url = f"https://opencre.org/cre/{cre_id}" if cre_id else "" + if url and url not in cites: + cites.append(url) + out["owasp_agent"] = {**meta, "citations": cites} + return out + + +def encode_citation_urls(response: Any) -> Any: + """Percent-encode spaces in citation URLs (repo file names like "4. Train your people.md"). + + A raw space breaks markdown links and makes the reference unclickable. Elementwise, so the + ``[n]`` marker to citation positions do not change. + """ + if not isinstance(response, dict): + return response + meta = response.get("owasp_agent") + if not isinstance(meta, dict) or not isinstance(meta.get("citations"), list): + return response + cites = [ + c.replace(" ", "%20") if isinstance(c, str) else c for c in meta["citations"] + ] + if cites == meta["citations"]: + return response + return {**response, "owasp_agent": {**meta, "citations": cites}} + + +def route_question( + prompt: str, history: Optional[List[Dict[str, str]]] = None +) -> Optional[Dict[str, Any]]: + """The OWASP agent's response for ``prompt``, or ``None`` when it does not handle it. + + Raises :class:`OwaspAgentUnavailable` when the agent crashes on a community/meta question + (we must not answer those from CRE retrieval); a crash on a CRE-shaped question returns + ``None`` so the caller can use the CRE path. + """ + from application.utils.owasp_agent import OwaspAgentRouter, is_owasp_agent_enabled + + if not is_owasp_agent_enabled(): + return None + try: + response = OwaspAgentRouter().handle(prompt or "", history=history or None) + if not response: + return response + return encode_citation_urls(attach_closest_cre(prompt or "", response)) + except Exception: # noqa: BLE001 + logger.exception("owasp agent router failed") + try: + from application.utils.owasp_agent.router import classify_intent + + if classify_intent(prompt or "") != "cre_normative": + raise OwaspAgentUnavailable() from None + except OwaspAgentUnavailable: + raise + except Exception: # noqa: BLE001 + logger.exception("owasp agent classify_intent also failed") + raise OwaspAgentUnavailable() from None + return None diff --git a/application/utils/rate_limit.py b/application/utils/rate_limit.py new file mode 100644 index 000000000..9bb899414 --- /dev/null +++ b/application/utils/rate_limit.py @@ -0,0 +1,155 @@ +"""Per-account question limits for the chatbot and MCP ``ask`` (each question costs model credits). + +Fixed windows (minute, hour, day) counted per signed-in email, shared by every web and MCP process through +Redis (falls back to in-process counters when Redis is unavailable). Admins get a multiple of the user limits. + +Defaults (per account): 5 per minute, 40 per hour, 150 per day; admins x4. +Override with RATE_LIMIT_PER_MINUTE / _PER_HOUR / _PER_DAY, RATE_LIMIT_ADMIN_MULTIPLIER; exempt accounts with +RATE_LIMIT_EXEMPT_EMAILS=a@x.org,b@y.org; disable with RATE_LIMIT_ENABLED=0. +""" + +from __future__ import annotations + +import logging +import os +import threading +import time +from typing import Dict, List, Optional, Tuple + +logger = logging.getLogger(__name__) + +WINDOWS: List[Tuple[str, int, str, int]] = [ + ("minute", 60, "RATE_LIMIT_PER_MINUTE", 5), + ("hour", 3600, "RATE_LIMIT_PER_HOUR", 40), + ("day", 86400, "RATE_LIMIT_PER_DAY", 150), +] + + +class RateLimited(Exception): + def __init__(self, window: str, limit: int, retry_after: int) -> None: + self.window, self.limit, self.retry_after = window, limit, retry_after + super().__init__(self.message()) + + def message(self) -> str: + wait = max(1, self.retry_after) + when = f"{wait} seconds" if wait < 120 else f"{(wait + 59) // 60} minutes" + return f"You have reached the limit of {self.limit} questions per {self.window}. Please try again in about {when}." + + +def enabled() -> bool: + return (os.environ.get("RATE_LIMIT_ENABLED", "1") or "1").strip().lower() not in ( + "0", + "false", + "no", + "off", + ) + + +def _limit(env_name: str, default: int) -> int: + try: + return max(1, int(os.environ.get(env_name, "") or default)) + except ValueError: + return default + + +def limits_for(role: Optional[str]) -> List[Tuple[str, int, int]]: + """[(window name, seconds, max questions)] for this role.""" + from application.utils import roles + + try: + mult = max(1, int(os.environ.get("RATE_LIMIT_ADMIN_MULTIPLIER", "") or 4)) + except ValueError: + mult = 4 + factor = mult if role == roles.ADMIN else 1 + return [ + (name, secs, _limit(env, default) * factor) + for name, secs, env, default in WINDOWS + ] + + +_mem: Dict[str, Tuple[int, float]] = {} +_mem_lock = threading.Lock() +_redis = {"conn": None, "tried": 0.0} + + +def _redis_conn(): + now = time.time() + if _redis["conn"] is not None or now - _redis["tried"] < 30: + return _redis["conn"] + _redis["tried"] = now + try: + from application.utils import redis as redis_utils + + conn = redis_utils.connect() + if conn is not None: + conn.ping() + _redis["conn"] = conn + except Exception: # noqa: BLE001 - fall back to in-process counters + _redis["conn"] = None + return _redis["conn"] + + +def _counts(keys: List[str]) -> List[int]: + conn = _redis_conn() + if conn is not None: + try: + return [int(v or 0) for v in conn.mget(keys)] + except Exception: # noqa: BLE001 + _redis["conn"] = None + now = time.time() + with _mem_lock: + return [(_mem[k][0] if k in _mem and _mem[k][1] > now else 0) for k in keys] + + +def _incr(keys: List[Tuple[str, int]]) -> None: + conn = _redis_conn() + if conn is not None: + try: + pipe = conn.pipeline() + for key, secs in keys: + pipe.incr(key) + pipe.expire(key, secs + 5) + pipe.execute() + return + except Exception: # noqa: BLE001 + _redis["conn"] = None + now = time.time() + with _mem_lock: + for key, secs in keys: + count, expires = _mem.get(key, (0, 0.0)) + _mem[key] = ((count if expires > now else 0) + 1, now + secs + 5) + + +def check_and_count( + account: Optional[str], role: Optional[str], scope: str = "ask" +) -> None: + """Raise :class:`RateLimited` when the account is over a limit; otherwise count this question.""" + if not enabled() or not (account or "").strip(): + return + who = account.strip().lower() + exempt = { + e.strip().lower() + for e in (os.environ.get("RATE_LIMIT_EXEMPT_EMAILS") or "").split(",") + if e.strip() + } + if ( + who in exempt + ): # e.g. the evaluation account that asks hundreds of questions in a row + return + now = time.time() + plan = [ + (f"rl:{scope}:{who}:{name}:{int(now // secs)}", name, secs, cap) + for name, secs, cap in limits_for(role) + ] + counts = _counts([p[0] for p in plan]) + for (key, name, secs, cap), used in zip(plan, counts): + if used >= cap: + raise RateLimited(name, cap, int(secs - (now % secs))) + _incr([(key, secs) for key, _, secs, _ in plan]) + + +def reset_memory() -> None: + """Test helper.""" + with _mem_lock: + _mem.clear() + _redis.update(conn=None, tried=time.time()) diff --git a/application/utils/roles.py b/application/utils/roles.py new file mode 100644 index 000000000..7644951f0 --- /dev/null +++ b/application/utils/roles.py @@ -0,0 +1,79 @@ +"""Who may do what: two roles shared by the web app and the MCP server. + +``user`` any account on an allowed domain (``LOGIN_ALLOWED_DOMAINS``, e.g. ``owasp.org``) may + use the chatbot, MCP ``ask`` and the read-only tools. +``admin`` the named accounts in ``ADMIN_EMAILS`` may also use the admin panel, manage every API + key, and call admin-only MCP tools. + +Roles are derived from the email on every request (never stored in a token), so removing someone +from ``ADMIN_EMAILS`` takes effect immediately. + +Backward compatible: while ``ADMIN_EMAILS`` is unset the old behaviour stays — ``LOGIN_ALLOWED_EMAILS`` +is the single allowlist and everyone on it is an admin; nobody else has any role. +""" + +from __future__ import annotations + +import os +from typing import Optional, Set + +USER = "user" +ADMIN = "admin" +_RANK = {USER: 1, ADMIN: 2} + + +def _csv(name: str) -> Set[str]: + return { + e.strip().lower() for e in (os.environ.get(name) or "").split(",") if e.strip() + } + + +def roles_enabled() -> bool: + """True once ``ADMIN_EMAILS`` is set (the opt-in switch for the two-role model).""" + return bool((os.environ.get("ADMIN_EMAILS") or "").strip()) + + +def _is_production() -> bool: + return (os.environ.get("FLASK_CONFIG") or "").strip().lower() == "production" + + +def admin_emails() -> Set[str]: + return _csv("ADMIN_EMAILS") if roles_enabled() else _csv("LOGIN_ALLOWED_EMAILS") + + +def _domain_allowed(email: str) -> bool: + domains = [ + d.strip().lower() + for d in (os.environ.get("LOGIN_ALLOWED_DOMAINS") or "").split(",") + if d.strip() + ] + if not domains: + return False # fail closed: a domain must be named explicitly + if domains == ["*"]: + return True + host = email.rsplit("@", 1)[-1] if "@" in email else "" + return any( + host == d.lstrip("@") or host.endswith("." + d.lstrip("@")) for d in domains + ) + + +def role_for(email: Optional[str]) -> Optional[str]: + """``admin``, ``user`` or ``None`` (no access) for this account.""" + addr = (email or "").strip().lower() + if not roles_enabled(): + # Legacy single allowlist: on it = admin; an empty list allows anyone outside production + # (even a session without an email, as before) and nobody in production. + allow = admin_emails() + if not allow: + return None if _is_production() else ADMIN + return ADMIN if addr in allow else None + if not addr: + return None + if addr in admin_emails(): + return ADMIN + return USER if _domain_allowed(addr) else None + + +def has_role(email: Optional[str], needed: str) -> bool: + role = role_for(email) + return role is not None and _RANK[role] >= _RANK[needed] diff --git a/application/web/admin_panel_routes.py b/application/web/admin_panel_routes.py index 037782ed5..e01a4b639 100644 --- a/application/web/admin_panel_routes.py +++ b/application/web/admin_panel_routes.py @@ -343,6 +343,42 @@ def admin_ingest_status() -> Any: logger.exception("ingest status failed") return _err(500, f"ingest status failed: {exc}") + @bp.route("/admin/authority", methods=["GET"]) + @login_required + @imports_enabled + def admin_authority_get() -> Any: + from application import sqla + from application.utils.comprehension import authority, retrieve + + try: + with sqla.session.get_bind().connect() as conn: + slugs = retrieve._known_slugs(conn) + return jsonify({"projects": authority.project_table(slugs)}) + except Exception as exc: + logger.exception("authority get failed") + return _err(500, f"authority failed: {exc}") + + @bp.route("/admin/authority", methods=["PUT"]) + @login_required + @imports_enabled + def admin_authority_put() -> Any: + from application.utils.comprehension import authority + + body = request.get_json(silent=True) or {} + try: + current = dict(authority.weights(fresh=True)) + current.update(body.get("weights") or {}) + if not isinstance(body.get("weights"), dict): + return _err(400, "weights object required") + actor = session.get("email") or getattr(g, "api_key_email", None) + saved = authority.save(current, actor=actor) + except ValueError as exc: + return _err(400, str(exc)) + except Exception as exc: + logger.exception("authority update failed") + return _err(500, f"authority update failed: {exc}") + return jsonify({"weights": saved}) + @bp.route("/admin/schedule", methods=["GET"]) @login_required @imports_enabled @@ -378,6 +414,26 @@ def admin_schedule_put(job: str) -> Any: logger.exception("schedule update failed") return _err(500, f"schedule update failed: {exc}") + @bp.route("/admin/feedback", methods=["GET"]) + @login_required + @imports_enabled + def admin_feedback_list() -> Any: + """Rated answers, newest first; ``?rating=down&limit=50``. Source for eval sets.""" + from application.utils.answer_feedback import RATINGS, list_feedback + + rating = request.args.get("rating") + if rating and rating not in RATINGS: + return _err(400, f"rating must be one of: {', '.join(RATINGS)}") + try: + limit = int(request.args.get("limit", 100)) + except ValueError: + return _err(400, "limit must be an integer") + try: + return jsonify({"items": list_feedback(rating=rating, limit=limit)}) + except Exception as exc: + logger.exception("feedback list failed") + return _err(500, f"feedback list failed: {exc}") + @bp.route("/admin/oie/rq/status", methods=["GET"]) @login_required @imports_enabled diff --git a/application/web/openapi_registry.py b/application/web/openapi_registry.py index fae08699b..f31c23ddb 100644 --- a/application/web/openapi_registry.py +++ b/application/web/openapi_registry.py @@ -42,6 +42,10 @@ "/rest/v1/api-keys", "/rest/v1/api-keys/", "/rest/v1/api-keys//audit", + # Signed-in account endpoints (role + personal API keys); session auth, not public REST. + "/rest/v1/me", + "/rest/v1/my-api-keys", + "/rest/v1/my-api-keys/", "/api/capabilities", "/docs/faq.md", "/", diff --git a/application/web/web_main.py b/application/web/web_main.py index 896460373..282e08cb3 100644 --- a/application/web/web_main.py +++ b/application/web/web_main.py @@ -13,7 +13,7 @@ import pathlib import re import urllib.parse -from typing import Any +from typing import Any, Optional from rq import job, exceptions from rq import Queue @@ -1001,6 +1001,9 @@ def _auth_challenge(): ) +from application.utils import roles as account_roles + + def _is_production_config() -> bool: return (os.environ.get("FLASK_CONFIG") or "").strip().lower() == "production" @@ -1010,16 +1013,18 @@ def _login_emails_raw() -> str: def _login_email_allowed(email: str | None) -> bool: - """True when ``email`` is on LOGIN_ALLOWED_EMAILS. + """Admin-level gate: True when ``email`` has the admin role. - Dev/test: unset/empty allowlist allows anyone (local convenience). - ``FLASK_CONFIG=production``: empty allowlist fails closed. + Legacy mode (``ADMIN_EMAILS`` unset): the account is on LOGIN_ALLOWED_EMAILS; dev/test with an + empty allowlist allows anyone; ``FLASK_CONFIG=production`` with an empty allowlist fails closed. + Roles mode (``ADMIN_EMAILS`` set): the account is one of the named admins. """ - raw = _login_emails_raw() - if not raw: - return not _is_production_config() - allowed = {e.strip().lower() for e in raw.split(",") if e.strip()} - return (email or "").strip().lower() in allowed + return account_roles.has_role(email, account_roles.ADMIN) + + +def _user_email_allowed(email: str | None) -> bool: + """User-level gate: any account with a role (user or admin).""" + return account_roles.has_role(email, account_roles.USER) def _ensure_csrf_token() -> str: @@ -1059,15 +1064,21 @@ def _require_same_origin_admin() -> None: abort(403, description="Cross-origin admin mutation refused") -def login_required(f): +def _auth_gate(f, allowed, role_label): + """Shared body of ``login_required`` (admin) and ``user_required`` (any role).""" + @wraps(f) - def login_r(*args, **kwargs): + def gated(*args, **kwargs): if os.environ.get("NO_LOGIN") == "1": return f(*args, **kwargs) - if _is_production_config() and not _login_emails_raw(): + if ( + _is_production_config() + and not _login_emails_raw() + and not account_roles.roles_enabled() + ): abort( 403, - description="LOGIN_ALLOWED_EMAILS must be set when FLASK_CONFIG=production", + description="LOGIN_ALLOWED_EMAILS (or ADMIN_EMAILS) must be set when FLASK_CONFIG=production", ) if not _is_logged_in(): from application.utils.api_tokens import verify as api_verify @@ -1077,26 +1088,40 @@ def login_r(*args, **kwargs): identity = api_verify.verify_api_token(raw) if identity is None: abort(401, description="Invalid or expired API key") - if not _login_email_allowed(identity.email): + if not allowed(identity.email): abort( 403, - description="API key owner is not on the allowlist for this environment.", + description=f"API key owner does not have the {role_label} role in this environment.", ) api_verify.bind_api_key_identity(identity) else: return _auth_challenge() - # Defense in depth: drop sessions that are no longer allowlisted. + # Defense in depth: re-check the role on every request. email = session.get("email") or getattr(g, "api_key_email", None) - if not _login_email_allowed(email): - if "user_id" in session: - session.clear() + if not allowed(email): + if "user_id" in session and not _user_email_allowed(email): + session.clear() # no role at all any more: drop the session abort( 403, - description="Your Google account is not on the allowlist for this environment.", + description=( + "Your Google account is not on the allowlist for this environment." + if not _user_email_allowed(email) + else f"The {role_label} role is required for this action." + ), ) return f(*args, **kwargs) - return login_r + return gated + + +def login_required(f): + """Admin-level gate (named admins). Used by the admin panel and API-key management.""" + return _auth_gate(f, _login_email_allowed, "admin") + + +def user_required(f): + """Any account with a role: the chatbot and answer feedback.""" + return _auth_gate(f, _user_email_allowed, "user") def login_or_mcp_required(f): @@ -1140,10 +1165,10 @@ def _admin_google_session_required(f): def wrapped(*args, **kwargs): if getattr(g, "api_key_jti", None) or not session.get("user_id"): abort(403, description="API key management requires Google login") - if not _login_emails_raw(): + if not _login_emails_raw() and not account_roles.roles_enabled(): abort( 403, - description="LOGIN_ALLOWED_EMAILS must be set to mint API keys", + description="LOGIN_ALLOWED_EMAILS (or ADMIN_EMAILS) must be set to mint API keys", ) if not _login_email_allowed(session.get("email")): abort(403, description="Only allowlisted admins can manage API keys") @@ -1532,13 +1557,57 @@ def admin_import_run_apply(run_id: str) -> Any: @app.route("/rest/v1/completion", methods=["POST"]) -@login_required +@user_required def chat_cre() -> Any: message = request.get_json(force=True) if posthog: posthog.capture("chat_cre", "") database = db.Node_collection() + + from application.utils import rate_limit + + if os.environ.get("NO_LOGIN") != "1": + _who = session.get("email") or getattr(g, "api_key_email", None) + try: + rate_limit.check_and_count(_who, account_roles.role_for(_who), "ask") + except rate_limit.RateLimited as limited: + resp = jsonify({"error": limited.message()}) + resp.status_code = 429 + resp.headers["Retry-After"] = str(limited.retry_after) + return resp + + from application.utils.chat_history import split_prompt_and_history + + user_prompt, chat_history = split_prompt_and_history( + message if isinstance(message, dict) else {} + ) + # Same routing and answer pipeline as MCP `ask` (application.utils.owasp_agent.service); + # the web chat additionally passes its conversation history. + from application.utils.owasp_agent.service import ( + OwaspAgentUnavailable, + route_question, + ) + + try: + agent_response = route_question(user_prompt or "", chat_history) + except OwaspAgentUnavailable: + return ( + jsonify( + { + "error": ( + "OWASP metadata agent failed. " + "Not falling back to CRE RAG for community/meta questions." + ) + } + ), + 503, + ) + if agent_response is not None: + return jsonify(_with_answer_id(agent_response, user_prompt)) + + # Only questions the OWASP agent does not handle need the CRE retrieval handler. Building it costs + # 2-3 s per request (16 s in a fresh process), so it is created here and never for agent answers. # Lazy import to avoid loading heavy prompt/ML dependencies at web boot. try: from application.prompt_client import llm_error_utils, prompt_client @@ -1557,51 +1626,6 @@ def chat_cre() -> Any: 503, "Chat is unavailable on this deployment (ML/LLM dependencies not installed).", ) - - from application.utils.chat_history import split_prompt_and_history - - user_prompt, chat_history = split_prompt_and_history( - message if isinstance(message, dict) else {} - ) - # Experimental OIE OWASP metadata agent (feature-flagged). Logistics/meta - # questions must not be forced through CRE RAG — see owasp_agent router. - try: - from application.utils.owasp_agent import ( - OwaspAgentRouter, - is_owasp_agent_enabled, - ) - - if is_owasp_agent_enabled(): - agent_response = OwaspAgentRouter().handle( - user_prompt or "", history=chat_history - ) - if agent_response is not None: - return jsonify(agent_response) - except Exception: - logger.exception("owasp agent router failed") - # If this looks like a meta question, do not fall through to CRE RAG. - try: - from application.utils.owasp_agent.router import classify_intent - - if classify_intent(user_prompt or "") != "cre_normative": - return ( - jsonify( - { - "error": ( - "OWASP metadata agent failed. " - "Not falling back to CRE RAG for community/meta questions." - ) - } - ), - 503, - ) - except Exception: - logger.exception("owasp agent classify_intent also failed") - return ( - jsonify({"error": "OWASP metadata agent unavailable"}), - 503, - ) - try: response = prompt.generate_text(user_prompt, history=chat_history) except Exception as e: @@ -1622,8 +1646,72 @@ def chat_cre() -> Any: jsonify({"error": f"AI Service Error: {str(e)}"}), status_code, ) + response = _block_unverified_answer(response, user_prompt) response = _maybe_append_cre_followups(user_prompt or "", response) - return jsonify(response) + return jsonify(_with_answer_id(response, user_prompt)) + + +def _block_unverified_answer(response: Any, question: Optional[str]) -> Any: + """Do not show free-form model text that no source backs (it can be a confident hallucination). + + ``generate_text`` marks an answer ``accurate`` only when a CRE/standard node matched the question. + Otherwise the model answered from its own memory ("who started opencre" -> a made-up founder), so the + user gets an honest "couldn't find it" instead. CHAT_BLOCK_UNVERIFIED=0 restores the old behaviour. + """ + blocking = ( + os.environ.get("CHAT_BLOCK_UNVERIFIED", "1") or "1" + ).strip().lower() not in ( + "0", + "false", + "no", + "off", + ) + if not blocking or not isinstance(response, dict) or response.get("accurate"): + return response + from application.utils.comprehension.answer import ABSTAIN_MESSAGE + + logger.info( + "blocked an unverified chat answer for %r: %s", + (question or "")[:120], + str(response.get("response") or "")[:300].replace("\n", " "), + ) + return {**response, "response": ABSTAIN_MESSAGE, "table": "", "accurate": False} + + +def _with_answer_id(response: Any, question: Optional[str]) -> Any: + """Log a chat answer and add its ``answer_id`` so the UI can offer thumbs up/down.""" + if not isinstance(response, dict): + return response + from application.utils.answer_feedback import log_answer + + citations = (response.get("owasp_agent") or {}).get("citations") or [] + answer_id = log_answer( + source="chat", + question=question or "", + answer=str(response.get("response") or ""), + citations=citations, + ) + return {**response, "answer_id": answer_id} if answer_id else response + + +@app.route("/rest/v1/answer_feedback", methods=["POST"]) +@user_required +def answer_feedback_post() -> Any: + """Thumbs up/down for an answer: ``{"answer_id", "rating": "up"|"down", "comment"?}``.""" + from application.utils.answer_feedback import record_rating + + body = request.get_json(silent=True) or {} + try: + result = record_rating( + str(body.get("answer_id") or ""), + str(body.get("rating") or ""), + str(body["comment"]) if body.get("comment") else None, + ) + except ValueError as exc: + abort(400, str(exc)) + except KeyError: + abort(404, "unknown answer_id") + return jsonify(result) def _maybe_append_cre_followups(prompt: str, response): @@ -1770,8 +1858,19 @@ def auth_callback(): 401, description=f"You need an account with one of the following providers to access this functionality {allowed_domains}", ) - # Optional exact-email allowlist (staging): further restricts who may log in. - allowed_emails_raw = (os.environ.get("LOGIN_ALLOWED_EMAILS") or "").strip() + if account_roles.roles_enabled(): + # Two-role model: anyone on an allowed domain may log in; named admins get the admin role. + if account_roles.role_for(email) is None: + abort( + 403, + description="Your Google account has no access to this environment.", + ) + # Legacy exact-email allowlist (staging): further restricts who may log in. + allowed_emails_raw = ( + "" + if account_roles.roles_enabled() + else (os.environ.get("LOGIN_ALLOWED_EMAILS") or "").strip() + ) if allowed_emails_raw: allowed_emails = { e.strip().lower() for e in allowed_emails_raw.split(",") if e.strip() @@ -2174,6 +2273,93 @@ def delete_user_mcp_grant() -> Any: # abort(404) +def _own_key_session_required(f): + """Personal API keys: any role, Google session only (not another API key), CSRF on writes.""" + + @wraps(f) + @user_required + def wrapped(*args, **kwargs): + if getattr(g, "api_key_jti", None) or not session.get("user_id"): + abort(403, description="API key management requires Google login") + if os.environ.get("NO_LOGIN") != "1" and request.method in ( + "POST", + "PUT", + "PATCH", + "DELETE", + ): + _ensure_csrf_token() + _require_csrf() + return f(*args, **kwargs) + + return wrapped + + +@app.route("/rest/v1/me", methods=["GET"]) +@user_required +def whoami() -> Any: + """The signed-in account and its role, so clients know which features to show.""" + email = session.get("email") or getattr(g, "api_key_email", None) + return jsonify({"email": email, "role": account_roles.role_for(email) or "none"}) + + +@app.route("/rest/v1/my-api-keys", methods=["POST"]) +@_own_key_session_required +def create_my_api_key() -> Any: + """Mint a personal API key for MCP/CLI use (shown once). Any account with a role.""" + from application.utils.api_tokens import mint as api_mint + + body = request.get_json(silent=True) or {} + name = body.get("name") if isinstance(body, dict) else None + expires_days = body.get("expires_days") if isinstance(body, dict) else None + database = db.Node_collection() + user = _resolve_current_user(database) + if user is None: + abort(401, description="Not authenticated") + try: + minted = api_mint.mint_api_token( + user=user, + name=str(name or "api-key"), + expires_days=int(expires_days) if expires_days is not None else None, + ) + except (TypeError, ValueError) as exc: + abort(400, description=str(exc)) + return ( + jsonify( + { + "token": minted.token, + "jti": minted.jti, + "name": minted.name, + "expires_at": minted.expires_at.isoformat() + "Z", + } + ), + 201, + ) + + +@app.route("/rest/v1/my-api-keys", methods=["GET"]) +@_own_key_session_required +def list_my_api_keys() -> Any: + from application.utils.api_tokens import mint as api_mint + + user = _resolve_current_user(db.Node_collection()) + if user is None: + abort(401, description="Not authenticated") + return jsonify({"keys": api_mint.list_api_tokens(user_id=user.id)}) + + +@app.route("/rest/v1/my-api-keys/", methods=["DELETE"]) +@_own_key_session_required +def revoke_my_api_key(jti: str) -> Any: + from application.utils.api_tokens import mint as api_mint + + user = _resolve_current_user(db.Node_collection()) + if user is None: + abort(401, description="Not authenticated") + if not api_mint.revoke_api_token(jti, user_id=user.id): + abort(404, description="API key not found") + return ("", 204) + + @app.route("/rest/v1/api-keys", methods=["POST"]) @_admin_google_session_required def create_api_key() -> Any: diff --git a/deploy/staging/Dockerfile.hotpatch b/deploy/staging/Dockerfile.hotpatch new file mode 100644 index 000000000..a74ae8fee --- /dev/null +++ b/deploy/staging/Dockerfile.hotpatch @@ -0,0 +1,18 @@ +# Layer application code from a synced tree onto the existing staging image (skips pip), +# and build the frontend bundle here so the deployed UI always matches the code. +# Build context = repository root: docker build -f deploy/staging/Dockerfile.hotpatch -t opencre-staging-web:latest . +FROM node:lts AS frontend +WORKDIR /code +COPY package.json yarn.lock /code/ +RUN yarn install --frozen-lockfile +COPY tsconfig.json webpack.config.js webpack.prod.js /code/ +COPY application/frontend /code/application/frontend +RUN yarn build + +FROM opencre-staging-web:latest +RUN pip install --no-cache-dir -U "litellm>=1.83.7,<1.107" +COPY application /code/application +COPY --from=frontend /code/application/frontend/www /code/application/frontend/www +COPY cre.py /code/cre.py +COPY scripts /code/scripts +COPY Makefile /code/Makefile diff --git a/docs/oie/admin-ingest-max-repos.md b/docs/oie/admin-ingest-max-repos.md new file mode 100644 index 000000000..e75da2e04 --- /dev/null +++ b/docs/oie/admin-ingest-max-repos.md @@ -0,0 +1,17 @@ +# `max_repos` on the admin ingest endpoints + +`POST /admin/ingest/start` (and the OIE run helpers behind it) accept an optional `max_repos`: + +| Value | Meaning | +|---|---| +| omitted / `null` | the capped default (`DEFAULT_OIE_MAX_REPOS`) | +| `N` > 0 | harvest at most N repos | +| `0` | **no cap**: every repo in the selected target | +| negative / not an integer | HTTP 400 | + +Why `0` means "unlimited" and not "reject": a full-organisation harvest is a legitimate, deliberate +operation (see `test_ingest_skip_c1_and_unlimited_max_repos`), and it needs an explicit value so it +can never happen by leaving the field out. The default stays capped so an accidental call is cheap. + +Implementation: `admin_panel_routes.py` validates (`>= 0`), `service.py` only passes `--max-repos` +to the harvester when the value is `> 0`. diff --git a/migrations/versions/d1e2f3a4b5c6_add_answer_log.py b/migrations/versions/d1e2f3a4b5c6_add_answer_log.py new file mode 100644 index 000000000..e29d82bde --- /dev/null +++ b/migrations/versions/d1e2f3a4b5c6_add_answer_log.py @@ -0,0 +1,40 @@ +"""answer_log: every answer given, with an optional thumbs up/down + +Revision ID: d1e2f3a4b5c6 +Revises: e5f6a7b8c9d0 +""" + +from alembic import op +import sqlalchemy as sa + + +revision = "d1e2f3a4b5c6" +down_revision = "e5f6a7b8c9d0" +branch_labels = None +depends_on = None + + +def upgrade(): + op.create_table( + "answer_log", + sa.Column("id", sa.String(), primary_key=True), + sa.Column("created_at", sa.DateTime(), nullable=False), + sa.Column("source", sa.String(), nullable=False), + sa.Column("question", sa.Text(), nullable=False), + sa.Column("answer", sa.Text(), nullable=True), + sa.Column("citations", sa.Text(), nullable=True), + sa.Column("channel", sa.String(), nullable=True), + sa.Column("flags", sa.String(), nullable=True), + sa.Column("latency_ms", sa.Integer(), nullable=True), + sa.Column("rating", sa.String(), nullable=True), + sa.Column("comment", sa.Text(), nullable=True), + sa.Column("rated_at", sa.DateTime(), nullable=True), + ) + op.create_index("ix_answer_log_created_at", "answer_log", ["created_at"]) + op.create_index("ix_answer_log_rating", "answer_log", ["rating"]) + + +def downgrade(): + op.drop_index("ix_answer_log_rating", table_name="answer_log") + op.drop_index("ix_answer_log_created_at", table_name="answer_log") + op.drop_table("answer_log") diff --git a/requirements.txt b/requirements.txt index 246ac8bda..ce400af55 100644 --- a/requirements.txt +++ b/requirements.txt @@ -45,7 +45,11 @@ semver python-markdown-maker # chat (/rest/v1/completion) — embed prompt via LiteLLM, match with sklearn -litellm +# Track the latest LiteLLM, never anything below 1.83.7: that release fixed the 2026 security advisories +# (https://github.com/BerriAI/litellm/security/advisories), and PyPI releases 1.82.7 / 1.82.8 were malicious. +# Re-read the advisories when upgrading. +# 1.107.x reported malicious (2026-10-26): stay below it until a clean release is confirmed. +litellm>=1.83.7,<1.107 numpy scipy scikit-learn diff --git a/scripts/comprehend_standard.py b/scripts/comprehend_standard.py index 3491f2f04..416a9e676 100644 --- a/scripts/comprehend_standard.py +++ b/scripts/comprehend_standard.py @@ -60,6 +60,26 @@ def main() -> int: action="store_true", help="with --raw: also index repo markdown files the harvester skipped (git clone)", ) + ap.add_argument( + "--repos-file", help="file with one repository per line (OWASP/x or x)" + ) + ap.add_argument( + "--clone-only", + action="store_true", + help="with --raw: take the chunks from a shallow clone instead of the harvest (repos never harvested)", + ) + ap.add_argument( + "--wiki", action="store_true", help="with --raw: also index the GitHub wiki" + ) + ap.add_argument( + "--wiki-only", action="store_true", help="with --raw: only the GitHub wiki" + ) + ap.add_argument( + "--max-chunks", + type=int, + default=0, + help="with --raw: cap chunks per repository (0 = no cap)", + ) ap.add_argument( "--supplement-only", action="store_true", @@ -108,6 +128,12 @@ def _meter(kwargs, response, start, end): # litellm success callback with app.app_context(): database = db_mod.Node_collection() repos = list(a.repo) + if a.repos_file: + with open(a.repos_file, encoding="utf-8") as fh: + for line in fh: + name = line.strip() + if name: + repos.append(name if "/" in name else f"OWASP/{name}") if a.like: from sqlalchemy import func @@ -153,7 +179,29 @@ def _meter(kwargs, response, start, end): # litellm success callback for repo in repos: try: if a.raw: - ch = sources.chunks_for_repo(database.session, repo) + import tempfile + + clone_dir = ( + os.environ.get("CRE_COMPREHENSION_CLONE_DIR") + or tempfile.mkdtemp() + ) + if a.clone_only or a.wiki_only: + ch = [] + else: + ch = sources.chunks_for_repo(database.session, repo) + if a.clone_only: + ch = sources.supplement_chunks(repo, set(), clone_dir) + print(f"{repo}: {len(ch)} chunks from the clone", flush=True) + if a.wiki or a.wiki_only: + wiki = sources.wiki_chunks(repo, clone_dir) + print(f"{repo}: {len(wiki)} wiki chunks", flush=True) + ch = ch + wiki + if a.max_chunks and len(ch) > a.max_chunks: + print( + f"{repo}: CAPPED {len(ch)} -> {a.max_chunks} chunks", + flush=True, + ) + ch = ch[: a.max_chunks] if a.supplement: import tempfile @@ -197,6 +245,15 @@ def _meter(kwargs, response, start, end): # litellm success callback workers=a.workers, dry_run=a.dry_run, enriched=enriched, + url_for=( + ( + lambda p, r=repo: ( + f"https://github.com/{r}/wiki" + if p.startswith("wiki/") + else f"https://github.com/{r}/blob/HEAD/{p}" + ) + ) + ), ) print( f"{repo}: raw chunks {'planned' if a.dry_run else 'written'}: {n}", diff --git a/scripts/ops/multiturn_probe.py b/scripts/ops/multiturn_probe.py new file mode 100644 index 000000000..2a84e2654 --- /dev/null +++ b/scripts/ops/multiturn_probe.py @@ -0,0 +1,109 @@ +#!/usr/bin/env python3 +"""Live multi-turn check, run INSIDE a deployed web container (uses the real models, costs a few cents). + + docker exec -w /code env PYTHONPATH=/code python scripts/ops/multiturn_probe.py + +Each scenario asks a question, then a follow-up that depends on it, sending the conversation as history +exactly like the chat page does. A follow-up answer must (1) mention what it was asked about and (2) NOT +repeat the previous answer. Exit 1 when a scenario fails. +""" + +from __future__ import annotations + +import difflib +import json +import os +import re +import sys +import time +from typing import Dict, List + +SCENARIOS: List[Dict] = [ + { + "name": "agentic AI risks", + "first": "Does the OWASP AI Exchange cover agentic AI?", + "follow_up": "What are some key risks", + "expect_any": ["risk", "threat", "attack", "injection"], + }, + { + "name": "SAMM levels", + "first": "What is OWASP SAMM?", + "follow_up": "How many levels does it define per practice?", + "expect_any": ["three", "3 "], + }, + { + "name": "membership by country", + "first": "What does OWASP membership cost in Switzerland?", + "follow_up": "and in Morocco?", + "expect_any": ["morocco"], + }, + { + "name": "project purpose", + "first": "Who leads the OWASP Developer Guide project?", + "follow_up": "What is its purpose?", + "expect_any": ["guide", "developer", "security"], + }, +] + + +def _norm(text: str) -> str: + return re.sub(r"\s+", " ", text or "").lower() + + +def evaluate(first_answer: str, second_answer: str, expect_any: List[str]) -> List[str]: + problems = [] + sim = difflib.SequenceMatcher( + None, _norm(first_answer), _norm(second_answer) + ).ratio() + if sim > 0.6: + problems.append(f"follow-up repeats the previous answer (similarity {sim:.2f})") + if not any(k in _norm(second_answer) for k in expect_any): + problems.append(f"follow-up answer mentions none of {expect_any}") + return problems + + +def main() -> int: + os.environ["NO_LOGIN"] = "1" + from application import create_app + + app = create_app(mode="production") + client = app.test_client() + failed = 0 + for sc in SCENARIOS: + t = time.time() + first = ( + client.post("/rest/v1/completion", json={"prompt": sc["first"]}).get_json() + or {} + ) + a1 = first.get("response") or first.get("error") or "" + history = [ + {"role": "user", "content": sc["first"]}, + {"role": "assistant", "content": a1}, + ] + second = ( + client.post( + "/rest/v1/completion", + json={"prompt": sc["follow_up"], "messages": history}, + ).get_json() + or {} + ) + a2 = second.get("response") or second.get("error") or "" + problems = evaluate(a1, a2, sc["expect_any"]) + failed += bool(problems) + print( + json.dumps( + { + "scenario": sc["name"], + "ok": not problems, + "problems": problems, + "seconds": round(time.time() - t, 1), + "follow_up_answer": a2[:160].replace("\n", " "), + } + ), + flush=True, + ) + return 1 if failed else 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/scripts/sync_check_staging.sh b/scripts/sync_check_staging.sh new file mode 100755 index 000000000..50c25039f --- /dev/null +++ b/scripts/sync_check_staging.sh @@ -0,0 +1,13 @@ +#!/usr/bin/env bash +# Print the stack tip on GitHub vs the commit the staging web image was built from. +# Usage: STAGING_HOST=sg@100.74.143.62 ./scripts/sync_check_staging.sh [branch] +set -euo pipefail +BRANCH="${1:-feat/answer-feedback}" +HOST="${STAGING_HOST:?set STAGING_HOST, e.g. user@host}" +git fetch -q origin "$BRANCH" +TIP="$(git rev-parse --short origin/"$BRANCH")" +# Staging is rsynced from a worktree checked out at the tip; compare file content, not a SHA. +LOCAL="$(git ls-tree -r origin/"$BRANCH" --name-only -- application/utils/comprehension | xargs -I{} git show origin/"$BRANCH":{} | shasum | cut -c1-12)" +REMOTE="$(ssh -n -o BatchMode=yes "$HOST" 'cd /home/sg/Projects/OpenCRE && for f in $(find application/utils/comprehension -type f -name "*.py" | sort); do cat "$f"; done | shasum | cut -c1-12')" +echo "GitHub $BRANCH tip: $TIP" +echo "comprehension code hash github=$LOCAL staging=$REMOTE" diff --git a/scripts/talks_catalogue_loader.py b/scripts/talks_catalogue_loader.py deleted file mode 100644 index 60315ffed..000000000 --- a/scripts/talks_catalogue_loader.py +++ /dev/null @@ -1,56 +0,0 @@ -#!/usr/bin/env python3 -"""Index a talk catalogue (one video per line: id|YYYYMMDD|seconds|title) as retrievable documents. - - python scripts/talks_catalogue_loader.py catalogue.txt [--repo OWASP/youtube-talks] [--dry-run] - -The catalogue comes from yt-dlp --flat-playlist --extractor-args youtubetab:approximate_date ---print '%(id)s|%(upload_date)s|%(duration)s|%(title)s' /videos . Dates from that listing are -approximate (YouTube only shows "N years ago"), so documents say so. Captions are not fetched. -""" - -from __future__ import annotations - -import argparse -import os -import sys - - -def main() -> int: - ap = argparse.ArgumentParser(description=__doc__) - ap.add_argument("catalogue") - ap.add_argument("--repo", default="OWASP/youtube-talks") - ap.add_argument("--dry-run", action="store_true") - ap.add_argument("--workers", type=int, default=8) - a = ap.parse_args() - - from application import create_app - from application.database import db as db_mod - from application.utils.comprehension import store - from application.utils.comprehension.shared import shared_handler - from application.utils.comprehension.talks_source import ( - catalogue_chunks, - parse_catalogue, - ) - - with open(a.catalogue, encoding="utf-8") as fh: - talks = parse_catalogue(fh.read()) - chunks = catalogue_chunks(talks) - print(f"talks: {len(talks)} chunks: {len(chunks)}", flush=True) - app = create_app(mode=os.environ.get("FLASK_CONFIG", "production")) - with app.app_context(): - database = db_mod.Node_collection() - n = store.write_raw( - database, - a.repo, - chunks, - shared_handler(database).get_text_embeddings, - workers=a.workers, - dry_run=a.dry_run, - url_for=lambda _p: "https://www.youtube.com/@OWASPGLOBAL/videos", - ) - print(f"catalogue chunks {'planned' if a.dry_run else 'written'}: {n}", flush=True) - return 0 - - -if __name__ == "__main__": - sys.exit(main())