From 9af0dfe336f6abd01db8c6c668a5a82e050c1f2e Mon Sep 17 00:00:00 2001 From: Robby Manihani Date: Mon, 18 May 2026 17:46:12 -0700 Subject: [PATCH] =?UTF-8?q?feat:=20Deep=20Research=20=E2=80=94=20personal?= =?UTF-8?q?=20deep=20research=20over=20Gmail=20with=20hybrid=20retrieval?= =?UTF-8?q?=20and=20agentic=20synthesis=20(#354)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- .github/workflows/ci.yml | 2 + .gitignore | 3 + frontend/src/components/Chat/ChatArea.tsx | 31 +- frontend/src/components/Chat/InputArea.tsx | 243 ++++++- .../src/components/Chat/MessageBubble.tsx | 38 +- .../src/components/Chat/ResearchTimeline.tsx | 236 ++++++ frontend/src/components/Chat/SystemPanel.tsx | 7 +- frontend/src/components/Chat/XRayFooter.tsx | 19 +- frontend/src/components/Sidebar/Sidebar.tsx | 10 +- frontend/src/index.css | 46 ++ frontend/src/lib/api.ts | 26 +- frontend/src/lib/rehype-citations.ts | 203 ++++++ frontend/src/lib/sse.ts | 51 +- frontend/src/lib/store.ts | 44 ++ frontend/src/pages/DataSourcesPage.tsx | 338 ++++++--- frontend/src/types/connectors.ts | 23 +- frontend/src/types/index.ts | 57 ++ frontend/vite.config.ts | 3 +- pyproject.toml | 5 + src/openjarvis/agents/executor.py | 26 +- src/openjarvis/agents/manager.py | 30 + src/openjarvis/agents/research_loop.py | 683 ++++++++++++++++++ src/openjarvis/analytics/identity.py | 15 +- src/openjarvis/cli/__init__.py | 10 +- src/openjarvis/cli/_version_check.py | 9 +- src/openjarvis/cli/ask.py | 215 ++++++ src/openjarvis/connectors/_stubs.py | 9 + src/openjarvis/connectors/chunker.py | 344 ++++++--- src/openjarvis/connectors/embeddings.py | 156 ++++ src/openjarvis/connectors/gcalendar.py | 18 +- src/openjarvis/connectors/gcontacts.py | 9 +- src/openjarvis/connectors/gdrive.py | 16 +- src/openjarvis/connectors/gmail.py | 182 ++++- src/openjarvis/connectors/gmail_imap.py | 30 +- src/openjarvis/connectors/google_auth.py | 130 ++++ src/openjarvis/connectors/google_tasks.py | 16 +- src/openjarvis/connectors/hybrid_search.py | 471 ++++++++++++ src/openjarvis/connectors/pipeline.py | 95 ++- src/openjarvis/connectors/store.py | 131 +++- src/openjarvis/engine/ollama.py | 8 + src/openjarvis/server/agent_manager_routes.py | 6 +- src/openjarvis/server/app.py | 2 + src/openjarvis/server/connectors_router.py | 259 ++++--- src/openjarvis/server/research_router.py | 478 ++++++++++++ tests/agents/test_research_loop.py | 186 +++++ tests/connectors/test_gmail.py | 496 ++++++++++++- tests/connectors/test_integration.py | 2 +- tests/connectors/test_pipeline.py | 171 +++++ tests/connectors/test_store.py | 106 +++ uv.lock | 2 + 50 files changed, 5253 insertions(+), 443 deletions(-) create mode 100644 frontend/src/components/Chat/ResearchTimeline.tsx create mode 100644 frontend/src/lib/rehype-citations.ts create mode 100644 src/openjarvis/agents/research_loop.py create mode 100644 src/openjarvis/connectors/embeddings.py create mode 100644 src/openjarvis/connectors/google_auth.py create mode 100644 src/openjarvis/connectors/hybrid_search.py create mode 100644 src/openjarvis/server/research_router.py create mode 100644 tests/agents/test_research_loop.py diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 6f3ce052..dd1ccecb 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -72,6 +72,8 @@ jobs: - name: Upload coverage report if: always() + continue-on-error: true + timeout-minutes: 5 uses: actions/upload-artifact@v4 with: name: coverage-xml diff --git a/.gitignore b/.gitignore index 92c34a67..4b7e2a7a 100644 --- a/.gitignore +++ b/.gitignore @@ -99,6 +99,9 @@ src/openjarvis/channels/whatsapp_baileys_bridge/node_modules/ # SQLite in-memory artifacts :memory: +# Dogfood reports (regenerated locally; not for VCS) +dogfood_report*.md + # Second Repos Inline/ scratch/ diff --git a/frontend/src/components/Chat/ChatArea.tsx b/frontend/src/components/Chat/ChatArea.tsx index 7bdd93a4..5df8b08c 100644 --- a/frontend/src/components/Chat/ChatArea.tsx +++ b/frontend/src/components/Chat/ChatArea.tsx @@ -146,14 +146,29 @@ export function ChatArea() { ) : (
- {messages.map((msg) => ( - - ))} - {streamState.isStreaming && streamState.content === '' && ( -
- -
- )} + {messages.map((msg, i) => { + const isLastAssistant = + i === messages.length - 1 && msg.role === 'assistant'; + return ( + + ); + })} + {(() => { + if (!streamState.isStreaming || streamState.content !== '') return null; + // For research messages the ResearchTimeline handles its own + // pre-content loading state — suppress the generic dots. + const last = messages[messages.length - 1]; + if (last?.role === 'assistant' && last.isResearch) return null; + return ( +
+ +
+ ); + })()}
)} diff --git a/frontend/src/components/Chat/InputArea.tsx b/frontend/src/components/Chat/InputArea.tsx index 13cdce9d..dd841b25 100644 --- a/frontend/src/components/Chat/InputArea.tsx +++ b/frontend/src/components/Chat/InputArea.tsx @@ -1,11 +1,77 @@ import { useState, useRef, useCallback, useEffect } from 'react'; -import { Send, Square, Paperclip } from 'lucide-react'; +import { Send, Square, Paperclip, Search } from 'lucide-react'; import { useAppStore, generateId } from '../../lib/store'; -import { streamChat } from '../../lib/sse'; +import { streamChat, streamResearch } from '../../lib/sse'; import { fetchSavings, getBase } from '../../lib/api'; +import { listConnectors, getSyncStatus } from '../../lib/connectors-api'; import { MicButton } from './MicButton'; import { useSpeech } from '../../hooks/useSpeech'; -import type { ChatMessage, ToolCallInfo, TokenUsage, MessageTelemetry } from '../../types'; +import type { + ChatMessage, + MessageTelemetry, + ResearchSearchTrace, + ResearchSource, + TokenUsage, + ToolCallInfo, +} from '../../types'; + +// While Deep Research is toggled on, poll connected sources for sync +// progress so we can surface "Searching over N items — sync in progress" +// next to the toggle. Polling is gated on `enabled` so toggling DR off +// stops the network chatter immediately. +function useResearchCorpusSync(enabled: boolean): { + syncing: boolean; + itemsSynced: number; +} { + const [state, setState] = useState({ syncing: false, itemsSynced: 0 }); + + useEffect(() => { + if (!enabled) { + setState({ syncing: false, itemsSynced: 0 }); + return; + } + let cancelled = false; + + const poll = async () => { + try { + const list = await listConnectors(); + const connected = list.filter((c) => c.connected); + if (connected.length === 0) { + if (!cancelled) setState({ syncing: false, itemsSynced: 0 }); + return; + } + const results = await Promise.all( + connected.map(async (c) => { + try { + return await getSyncStatus(c.connector_id); + } catch { + return null; + } + }), + ); + let syncing = false; + let itemsSynced = 0; + for (const r of results) { + if (!r) continue; + if (r.state === 'syncing') syncing = true; + itemsSynced += r.items_synced ?? 0; + } + if (!cancelled) setState({ syncing, itemsSynced }); + } catch { + // Network blip — leave previous state intact. + } + }; + + poll(); + const interval = setInterval(poll, 5000); + return () => { + cancelled = true; + clearInterval(interval); + }; + }, [enabled]); + + return state; +} export function InputArea() { const [input, setInput] = useState(''); @@ -26,6 +92,9 @@ export function InputArea() { const setStreamState = useAppStore((s) => s.setStreamState); const resetStream = useAppStore((s) => s.resetStream); const modelLoading = useAppStore((s) => s.modelLoading); + const deepResearch = useAppStore((s) => s.deepResearch); + const setDeepResearch = useAppStore((s) => s.setDeepResearch); + const corpusSync = useResearchCorpusSync(deepResearch); const { state: speechState, available: speechAvailable, startRecording, stopRecording } = useSpeech(); @@ -114,6 +183,7 @@ export function InputArea() { role: 'assistant', content: '', timestamp: Date.now(), + isResearch: deepResearch || undefined, }; addMessage(convId, assistantMsg); @@ -131,12 +201,16 @@ export function InputArea() { let usage: TokenUsage | undefined; let complexity: { score: number; tier: string; suggested_max_tokens: number } | undefined; const toolCalls: ToolCallInfo[] = []; + const researchTraces: ResearchSearchTrace[] = []; + const researchSourcesByRef = new Map(); + const flushSources = () => + Array.from(researchSourcesByRef.values()).sort((a, b) => a.ref - b.ref); let lastFlush = 0; let ttftMs: number | undefined; setStreamState({ isStreaming: true, - phase: 'Generating...', + phase: deepResearch ? 'Researching...' : 'Generating...', elapsedMs: 0, activeToolCalls: [], content: '', @@ -145,10 +219,116 @@ export function InputArea() { timestamp: Date.now(), level: 'info', category: 'chat', - message: `Request: "${content.slice(0, 80)}${content.length > 80 ? '...' : ''}" → ${selectedModel}`, + message: deepResearch + ? `Research: "${content.slice(0, 80)}${content.length > 80 ? '...' : ''}"` + : `Request: "${content.slice(0, 80)}${content.length > 80 ? '...' : ''}" → ${selectedModel}`, }); try { + if (deepResearch) { + for await (const ev of streamResearch(content, controller.signal)) { + if (ev.type === 'search_call') { + const trace: ResearchSearchTrace = { + id: generateId(), + query: ev.arguments?.query ?? '', + person: ev.arguments?.person, + timeRange: ev.arguments?.time_range, + status: 'pending', + }; + researchTraces.push(trace); + setStreamState({ phase: `Searching: ${trace.query}` }); + updateLastAssistant( + convId, + accumulatedContent, + undefined, + undefined, + undefined, + undefined, + [...researchTraces], + flushSources(), + ); + useAppStore.getState().addLogEntry({ + timestamp: Date.now(), + level: 'info', + category: 'tool', + message: `Search: "${trace.query}"${trace.person ? ` (person: ${trace.person})` : ''}`, + }); + } else if (ev.type === 'search_result') { + const pending = [...researchTraces].reverse().find((t) => t.status === 'pending'); + if (pending) { + pending.status = 'complete'; + pending.numHits = ev.num_hits; + pending.topTitles = ev.top_titles; + } + if (ev.sources) { + for (const src of ev.sources) { + if (src && typeof src.ref === 'number' && !researchSourcesByRef.has(src.ref)) { + researchSourcesByRef.set(src.ref, src); + } + } + } + updateLastAssistant( + convId, + accumulatedContent, + undefined, + undefined, + undefined, + undefined, + [...researchTraces], + flushSources(), + ); + } else if (ev.type === 'synthesis') { + if (!ttftMs) ttftMs = Date.now() - startTime; + accumulatedContent += ev.text; + setStreamState({ content: accumulatedContent, phase: '' }); + const now = Date.now(); + if (now - lastFlush >= 80) { + updateLastAssistant( + convId, + accumulatedContent, + undefined, + undefined, + undefined, + undefined, + [...researchTraces], + flushSources(), + ); + lastFlush = now; + } + } else if (ev.type === 'system_metrics') { + // Live GPU sample — feed straight to the System panel so Power + // (W) and Energy (kJ) tick up in real time as the agent runs. + useAppStore.getState().setLiveEnergy({ + power_w: ev.power_w, + energy_j: ev.energy_j, + duration_s: ev.duration_s, + }); + } else if (ev.type === 'done') { + if (ev.usage) { + usage = { + prompt_tokens: ev.usage.prompt_tokens ?? 0, + completion_tokens: ev.usage.completion_tokens ?? 0, + total_tokens: + ev.usage.total_tokens ?? + (ev.usage.prompt_tokens ?? 0) + + (ev.usage.completion_tokens ?? 0), + }; + // Optimistically roll this research turn into the session + // counters so the Session panel updates the moment the + // stream finishes, regardless of how /v1/savings aggregates + // research telemetry server-side. + useAppStore.getState().incrementSavings(usage); + } + // Hold the final live numbers visible for a beat so the panel + // doesn't flash to 0 between the SSE close and the next + // /v1/telemetry/energy poll picking up the persisted record. + window.setTimeout(() => { + useAppStore.getState().setLiveEnergy(null); + }, 1500); + break; + } + } + } else { for await (const sseEvent of streamChat( { model: selectedModel, messages: apiMessages, stream: true, temperature, max_tokens: maxTokens }, controller.signal, @@ -225,6 +405,7 @@ export function InputArea() { } catch {} } } + } } catch (err: any) { if (err.name === 'AbortError') { // User cancelled or model switch — keep whatever was accumulated @@ -238,6 +419,9 @@ export function InputArea() { message: `Stream error: ${errMsg}`, }); } + // If we tore out mid-research, make sure the live System panel + // numbers don't get stuck on the last sample. + useAppStore.getState().setLiveEnergy(null); } finally { if (!accumulatedContent) { accumulatedContent = 'No response was generated. Please try again.'; @@ -278,6 +462,8 @@ export function InputArea() { usage, telemetry, audioMeta, + researchTraces.length > 0 ? researchTraces : undefined, + researchSourcesByRef.size > 0 ? flushSources() : undefined, ); if (timerRef.current) { clearInterval(timerRef.current); @@ -290,9 +476,15 @@ export function InputArea() { }); abortRef.current = null; - fetchSavings() - .then((data) => useAppStore.getState().setSavings(data)) - .catch(() => {}); + // Research path updates session counters optimistically from the + // `done` event's usage payload — re-fetching here would overwrite + // it with a potentially stale snapshot if the server's research + // telemetry hasn't been merged into /v1/savings yet. + if (!deepResearch) { + fetchSavings() + .then((data) => useAppStore.getState().setSavings(data)) + .catch(() => {}); + } } }, [ input, @@ -304,6 +496,9 @@ export function InputArea() { updateLastAssistant, setStreamState, resetStream, + deepResearch, + temperature, + maxTokens, ]); const handleKeyDown = (e: React.KeyboardEvent) => { @@ -315,6 +510,38 @@ export function InputArea() { return (
+
+
+ +
+ {deepResearch && corpusSync.syncing && corpusSync.itemsSynced > 0 && ( +
+ Searching over{' '} + + {corpusSync.itemsSynced.toLocaleString()} + {' '} + items — sync in progress, results will improve as more data is indexed. +
+ )} +
stripThinkTags(message.content), [message.content]); + // Build a ref→source lookup once per render. Memoized so the rehype plugin + // identity stays stable until the source list actually changes. + const sourcesMap = useMemo(() => { + const m = new Map[number]>(); + for (const s of message.researchSources ?? []) { + if (typeof s.ref === 'number') m.set(s.ref, s); + } + return m; + }, [message.researchSources]); + + const rehypePlugins = useMemo(() => { + const base: any[] = [[rehypeHighlight, { detect: true }], rehypeKatex]; + if (sourcesMap.size > 0) base.push([rehypeCitations, { sources: sourcesMap }]); + return base; + }, [sourcesMap]); + return (
+ {/* Deep Research timeline (steps + status) */} + {(message.isResearch || (message.researchTraces && message.researchTraces.length > 0)) && ( + 0} + /> + )} + {/* Tool calls */} {message.toolCalls && message.toolCalls.length > 0 && (
@@ -140,7 +168,7 @@ export function MessageBubble({ message }: Props) {
- +
); } diff --git a/frontend/src/components/Chat/ResearchTimeline.tsx b/frontend/src/components/Chat/ResearchTimeline.tsx new file mode 100644 index 00000000..be66a8af --- /dev/null +++ b/frontend/src/components/Chat/ResearchTimeline.tsx @@ -0,0 +1,236 @@ +import { useEffect, useRef, useState } from 'react'; +import { ChevronDown, ChevronUp } from 'lucide-react'; +import type { ResearchSearchTrace, TimeRange } from '../../types'; + +interface Props { + traces: ResearchSearchTrace[]; + isLive: boolean; + hasContent: boolean; +} + +const SHORT_DATE = new Intl.DateTimeFormat('en-US', { + month: 'short', + day: 'numeric', +}); + +function fmtDate(iso: string): string | null { + const d = new Date(iso); + if (Number.isNaN(d.getTime())) return null; + return SHORT_DATE.format(d); +} + +function formatTimeRange(tr: TimeRange | string | undefined): string | null { + if (!tr) return null; + if (typeof tr === 'string') return tr.trim() || null; + const start = tr.start ? fmtDate(tr.start) : null; + const end = tr.end ? fmtDate(tr.end) : null; + if (start && end) return start === end ? start : `${start} – ${end}`; + if (start) return `after ${start}`; + if (end) return `before ${end}`; + return null; +} + +function summarizeTraces(traces: ResearchSearchTrace[]): string { + const n = traces.length; + const hits = traces.reduce((s, t) => s + (t.numHits ?? 0), 0); + const searchLabel = `${n} ${n === 1 ? 'search' : 'searches'}`; + if (hits === 0) return searchLabel; + return `${searchLabel} · ${hits} ${hits === 1 ? 'result' : 'results'}`; +} + +function StatusLine({ text }: { text: string }) { + return ( +
+ {text}… +
+ ); +} + +function TimelineStep({ + index, + trace, + isLive, +}: { + index: number; + trace: ResearchSearchTrace; + isLive: boolean; +}) { + const pending = trace.status === 'pending'; + const active = pending && isLive; + const meta: string[] = []; + if (trace.person) meta.push(`person: ${trace.person}`); + const formattedTime = formatTimeRange(trace.timeRange); + if (formattedTime) meta.push(`time: ${formattedTime}`); + + return ( +
+ {/* Step indicator dot, sitting on the vertical rail */} +
+ +
+ Search {index} +
+ +
+ “{trace.query}” +
+ + {meta.length > 0 && ( +
+ {meta.join(' · ')} +
+ )} + + {active ? ( +
+
+
+ ) : trace.numHits != null ? ( +
+ {trace.numHits} {trace.numHits === 1 ? 'result' : 'results'} +
+ ) : null} + + {trace.topTitles && trace.topTitles.length > 0 && ( +
+ {trace.topTitles.slice(0, 2).join(' · ')} +
+ )} +
+ ); +} + +export function ResearchTimeline({ traces, isLive, hasContent }: Props) { + const showAnalyzing = isLive && traces.length === 0 && !hasContent; + const allComplete = + traces.length > 0 && traces.every((t) => t.status === 'complete'); + const showSynthesizing = isLive && allComplete && !hasContent; + + // Auto-collapse the timeline the moment synthesis text begins streaming. + // Subsequent user toggles win — we only fire the auto-collapse once per + // false→true transition of hasContent. + const [expanded, setExpanded] = useState(true); + const prevHasContent = useRef(false); + useEffect(() => { + if (hasContent && !prevHasContent.current) { + setExpanded(false); + } + prevHasContent.current = hasContent; + }, [hasContent]); + + if (!showAnalyzing && traces.length === 0) return null; + + // No collapse affordance before any traces have arrived — just the + // analyzing status sitting alone. + if (traces.length === 0) { + return ( +
+
+
+ +
+
+ ); + } + + const summary = summarizeTraces(traces); + const Chevron = expanded ? ChevronUp : ChevronDown; + + return ( +
+ + + {expanded && ( +
+
+
+ {traces.map((t, i) => ( + + ))} + + {showSynthesizing && ( +
+
+ +
+ )} +
+
+ )} +
+ ); +} diff --git a/frontend/src/components/Chat/SystemPanel.tsx b/frontend/src/components/Chat/SystemPanel.tsx index e83b2e9b..0fc2a866 100644 --- a/frontend/src/components/Chat/SystemPanel.tsx +++ b/frontend/src/components/Chat/SystemPanel.tsx @@ -39,6 +39,7 @@ export function SystemPanel() { const toggleSystemPanel = useAppStore((s) => s.toggleSystemPanel); const optInEnabled = useAppStore((s) => s.optInEnabled); const setOptInModalOpen = useAppStore((s) => s.setOptInModalOpen); + const liveEnergy = useAppStore((s) => s.liveEnergy); const [energy, setEnergy] = useState(null); const [telemetry, setTelemetry] = useState(null); @@ -129,13 +130,15 @@ export function SystemPanel() {
diff --git a/frontend/src/components/Chat/XRayFooter.tsx b/frontend/src/components/Chat/XRayFooter.tsx index b13126dc..08fb3468 100644 --- a/frontend/src/components/Chat/XRayFooter.tsx +++ b/frontend/src/components/Chat/XRayFooter.tsx @@ -5,19 +5,26 @@ import type { TokenUsage, MessageTelemetry } from '../../types'; interface Props { usage?: TokenUsage; telemetry?: MessageTelemetry; + isResearch?: boolean; } function formatMs(ms: number): string { return ms < 1000 ? `${Math.round(ms)}ms` : `${(ms / 1000).toFixed(1)}s`; } -export function XRayFooter({ usage, telemetry }: Props) { +export function XRayFooter({ usage, telemetry, isResearch = false }: Props) { const [expanded, setExpanded] = useState(false); - // Build collapsed summary parts + // Build collapsed summary parts. For Deep Research responses we hide the + // chat-side engine/model (which doesn't reflect the planner that actually + // produced the answer) and label the response with the mode instead. const parts: string[] = []; - if (telemetry?.engine) parts.push(telemetry.engine); - if (telemetry?.model_id) parts.push(telemetry.model_id); + if (isResearch) { + parts.push('Deep Research'); + } else { + if (telemetry?.engine) parts.push(telemetry.engine); + if (telemetry?.model_id) parts.push(telemetry.model_id); + } if (telemetry?.complexity_tier) parts.push(telemetry.complexity_tier); if (telemetry?.total_ms) parts.push(formatMs(telemetry.total_ms)); if (usage && (usage.prompt_tokens || usage.completion_tokens)) { @@ -32,7 +39,9 @@ export function XRayFooter({ usage, telemetry }: Props) { // Build expanded rows const rows: Array<{ label: string; value: string; color?: string }> = []; - if (telemetry?.engine) { + if (isResearch) { + rows.push({ label: 'Mode', value: 'Deep Research' }); + } else if (telemetry?.engine) { const modelDetail = telemetry.model_id || ''; rows.push({ label: 'Engine', value: `${telemetry.engine}${modelDetail ? ` (${modelDetail})` : ''}` }); } diff --git a/frontend/src/components/Sidebar/Sidebar.tsx b/frontend/src/components/Sidebar/Sidebar.tsx index de518405..a4a5b670 100644 --- a/frontend/src/components/Sidebar/Sidebar.tsx +++ b/frontend/src/components/Sidebar/Sidebar.tsx @@ -33,6 +33,7 @@ export function Sidebar() { const serverInfo = useAppStore((s) => s.serverInfo); const setCommandPaletteOpen = useAppStore((s) => s.setCommandPaletteOpen); const modelLoading = useAppStore((s) => s.modelLoading); + const deepResearch = useAppStore((s) => s.deepResearch); const settings = useAppStore((s) => s.settings); const updateSettings = useAppStore((s) => s.updateSettings); @@ -143,8 +144,13 @@ export function Sidebar() { )}
- - {selectedModel || serverInfo?.model || 'Select model'} + + {deepResearch + ? 'Deep Research' + : selectedModel || serverInfo?.model || 'Select model'} {modelLoading && ( diff --git a/frontend/src/index.css b/frontend/src/index.css index b043b74a..a974c6e8 100644 --- a/frontend/src/index.css +++ b/frontend/src/index.css @@ -707,4 +707,50 @@ --color-foreground: var(--foreground); --color-background: var(--background); /* Keep project radius tokens — don't override with shadcn calc values */ +} + +/* Deep Research timeline — pending-step shimmer */ +@keyframes openjarvis-research-shimmer { + 0% { transform: translateX(-120%); } + 100% { transform: translateX(420%); } +} +.research-shimmer { + animation: openjarvis-research-shimmer 1.6s ease-in-out infinite; +} + +/* Sync progress pulse — fires on each items_synced bump so the + "12,450 / 99,840" label visibly breathes between polls. */ +@keyframes openjarvis-sync-bump { + 0% { opacity: 0.55; } + 100% { opacity: 1; } +} +.sync-bump { + animation: openjarvis-sync-bump 400ms ease-out; +} + +/* Deep Research inline citation pills — small rounded chips that flow + with the surrounding text and link out to the cited email. */ +.research-citation { + display: inline-block; + vertical-align: baseline; + margin: 0 1px; + padding: 0 6px; + min-width: 1.25em; + text-align: center; + font-size: 0.72em; + font-weight: 600; + line-height: 1.5; + color: var(--color-accent); + background: var(--color-accent-subtle); + border-radius: 9999px; + text-decoration: none; + cursor: pointer; + transition: background 150ms, color 150ms; + /* Keep the pill from inheriting prose link styles (underline, bold). */ + border-bottom: none; +} +.research-citation:hover { + color: var(--color-on-accent, #ffffff); + background: var(--color-accent); + text-decoration: none; } \ No newline at end of file diff --git a/frontend/src/lib/api.ts b/frontend/src/lib/api.ts index 32a4d173..b62d71b2 100644 --- a/frontend/src/lib/api.ts +++ b/frontend/src/lib/api.ts @@ -192,12 +192,26 @@ export async function checkHealth(): Promise { return false; } } - try { - const res = await fetch(`${getBase()}/health`); - return res.ok; - } catch { - return false; - } + // In the browser, hit /health relative to the page origin so the request + // flows through whatever path is already serving the SPA — the Vite + // proxy in dev, FastAPI's static mount in prod. This avoids the + // false-negative "Cannot reach backend" banner when getBase() points at + // an absolute URL the browser can't reach directly. + // + // If /health itself fails for any reason (proxy quirk, stale service + // worker, etc.) fall back to an arbitrary API endpoint we know the rest + // of the app polls successfully. If THAT also fails we genuinely can't + // reach the backend. + const probe = async (url: string): Promise => { + try { + const res = await fetch(url, { cache: 'no-store' }); + return res.ok; + } catch { + return false; + } + }; + if (await probe('/health')) return true; + return probe('/v1/connectors'); } export async function fetchEnergy(): Promise { diff --git a/frontend/src/lib/rehype-citations.ts b/frontend/src/lib/rehype-citations.ts new file mode 100644 index 00000000..18332d1a --- /dev/null +++ b/frontend/src/lib/rehype-citations.ts @@ -0,0 +1,203 @@ +import type { ResearchSource } from '../types'; + +interface HastNode { + type: string; + tagName?: string; + value?: string; + properties?: Record; + children?: HastNode[]; +} + +// Tags whose descendants should NOT be searched for citations. +// `` is handled specially below (markdown link to a single number is +// promoted to a pill); other links + code/pre/script/style are skipped. +const SKIP_TAGS = new Set(['code', 'pre', 'script', 'style']); + +// Matches `[N]` or `[N, M, ...]` — bracketed comma-separated digit lists. +// Whitespace inside the brackets is tolerated. An optional trailing +// whitespace run is consumed when followed by sentence punctuation, so we +// can render `… San Francisco [1].` without a stray space before the period. +const CITATION_RE = + /\[(\s*\d+(?:\s*,\s*\d+)*\s*)\](\s+(?=[.,;:!?]))?/g; + +const TOOLTIP_DATE = new Intl.DateTimeFormat('en-US', { + month: 'short', + day: 'numeric', + year: 'numeric', +}); + +function formatTooltipDate(iso: string | undefined): string | null { + if (!iso) return null; + const d = new Date(iso); + if (Number.isNaN(d.getTime())) return iso.trim() || null; + return TOOLTIP_DATE.format(d); +} + +function buildTooltip(src: ResearchSource): string { + const parts: string[] = []; + if (src.title) parts.push(src.title); + const meta: string[] = []; + if (src.sender) meta.push(src.sender); + const date = formatTooltipDate(src.date); + if (date) meta.push(date); + if (meta.length > 0) parts.push(meta.join(' · ')); + return parts.join('\n'); +} + +function buildPill(n: number, src: ResearchSource): HastNode { + return { + type: 'element', + tagName: 'a', + properties: { + href: src.url, + target: '_blank', + rel: 'noopener noreferrer', + className: ['research-citation'], + 'data-ref': String(n), + title: buildTooltip(src) || `Source ${n}`, + }, + children: [{ type: 'text', value: String(n) }], + }; +} + +function parseRefList(inner: string): number[] { + return inner + .split(',') + .map((s) => parseInt(s.trim(), 10)) + .filter((n) => !Number.isNaN(n)); +} + +/** + * Replace bracketed citations in the rendered markdown with small inline + * pill links to the underlying source. + * + * Patterns handled: + * • `[1]` — single ref + * • `[4, 7, 20]` — grouped refs → individual pills, no separators + * • `[1](https://…)` — markdown link whose text is just a number; the + * inline URL is discarded in favor of the source-map + * URL so all citations stay routed through Gmail. + * + * Citations whose ref is missing from the source map are preserved as + * literal text so we never silently lose information. + */ +export function rehypeCitations(options: { + sources: Map; +}) { + const { sources } = options; + if (sources.size === 0) { + return () => undefined; + } + + function transformText(text: string): HastNode[] | null { + if (!text.includes('[')) return null; + + const pieces: HastNode[] = []; + let lastIndex = 0; + let m: RegExpExecArray | null; + CITATION_RE.lastIndex = 0; + let changed = false; + + while ((m = CITATION_RE.exec(text)) !== null) { + const refs = parseRefList(m[1]); + const resolved = refs.map((n) => ({ n, src: sources.get(n) })); + const allMatched = + resolved.length > 0 && resolved.every((r) => r.src && r.src.url); + + if (m.index > lastIndex) { + pieces.push({ type: 'text', value: text.slice(lastIndex, m.index) }); + } + + if (allMatched) { + for (const r of resolved) { + pieces.push(buildPill(r.n, r.src!)); + } + changed = true; + } else { + // Conservative fallback: if any ref in this group is missing source + // data, keep the whole bracketed expression as the model wrote it. + pieces.push({ type: 'text', value: m[0] }); + } + + lastIndex = m.index + m[0].length; + } + + if (!changed) return null; + + if (lastIndex < text.length) { + pieces.push({ type: 'text', value: text.slice(lastIndex) }); + } + return pieces; + } + + function tryMarkdownCitationLink(node: HastNode): HastNode | null { + // `[1](url)` → 1. If the link text is just a number + // and we have a source for that ref, replace with our pill. + if ( + node.type !== 'element' || + node.tagName !== 'a' || + !node.children || + node.children.length !== 1 + ) { + return null; + } + const child = node.children[0]; + if (child.type !== 'text' || typeof child.value !== 'string') return null; + + const trimmed = child.value.trim(); + if (!/^\d+$/.test(trimmed)) return null; + + const n = parseInt(trimmed, 10); + const src = sources.get(n); + if (!src || !src.url) return null; + + return buildPill(n, src); + } + + function walk(node: HastNode): void { + if (!node.children || node.children.length === 0) return; + if (node.tagName && SKIP_TAGS.has(node.tagName)) return; + + const out: HastNode[] = []; + let changed = false; + + for (const child of node.children) { + // 1. Markdown citation link: [1](url) + const promoted = tryMarkdownCitationLink(child); + if (promoted) { + out.push(promoted); + changed = true; + continue; + } + + // 2. Other elements — recurse but don't transform inside /code/etc. + if (child.type === 'element') { + if (child.tagName !== 'a') walk(child); + out.push(child); + continue; + } + + // 3. Text node — look for [N] and [N, M, …] patterns. + if (child.type === 'text' && typeof child.value === 'string') { + const replaced = transformText(child.value); + if (replaced) { + out.push(...replaced); + changed = true; + } else { + out.push(child); + } + continue; + } + + out.push(child); + } + + if (changed) { + node.children = out; + } + } + + return (tree: HastNode) => { + walk(tree); + }; +} diff --git a/frontend/src/lib/sse.ts b/frontend/src/lib/sse.ts index d00fae92..ffe7baff 100644 --- a/frontend/src/lib/sse.ts +++ b/frontend/src/lib/sse.ts @@ -1,4 +1,4 @@ -import type { SSEEvent } from '../types'; +import type { ResearchEvent, SSEEvent } from '../types'; import { getBase } from './api'; export interface ChatRequest { @@ -57,3 +57,52 @@ export async function* streamChat( reader.releaseLock(); } } + +export async function* streamResearch( + query: string, + signal?: AbortSignal, +): AsyncGenerator { + // /api/research is mounted at the server root — strip any trailing /v1 + // from the base so configurations like "http://host:8000/v1" still resolve. + const base = getBase().replace(/\/v1\/?$/, ''); + const response = await fetch(`${base}/api/research`, { + method: 'POST', + headers: { 'Content-Type': 'application/json' }, + body: JSON.stringify({ query }), + signal, + }); + + if (!response.ok) { + throw new Error(`Research request failed: ${response.status}`); + } + + const reader = response.body!.getReader(); + const decoder = new TextDecoder(); + let buffer = ''; + + try { + while (true) { + const { done, value } = await reader.read(); + if (done) break; + + buffer += decoder.decode(value, { stream: true }); + const lines = buffer.split('\n'); + buffer = lines.pop() || ''; + + for (const line of lines) { + if (!line.startsWith('data: ')) continue; + const data = line.slice(6); + if (data === '[DONE]') return; + try { + const parsed = JSON.parse(data) as ResearchEvent; + yield parsed; + if (parsed.type === 'done') return; + } catch { + // skip malformed chunks + } + } + } + } finally { + reader.releaseLock(); + } +} diff --git a/frontend/src/lib/store.ts b/frontend/src/lib/store.ts index 2b52219a..37a7d84c 100644 --- a/frontend/src/lib/store.ts +++ b/frontend/src/lib/store.ts @@ -2,9 +2,12 @@ import { create } from 'zustand'; import type { Conversation, ChatMessage, + LiveEnergyMetrics, LogEntry, ModelInfo, MessageTelemetry, + ResearchSearchTrace, + ResearchSource, SavingsData, ServerInfo, StreamState, @@ -158,16 +161,29 @@ interface AppState { usage?: TokenUsage, telemetry?: MessageTelemetry, audio?: { url: string }, + researchTraces?: ResearchSearchTrace[], + researchSources?: ResearchSource[], ) => void; setStreamState: (state: Partial) => void; resetStream: () => void; + // Deep Research toggle + deepResearch: boolean; + setDeepResearch: (on: boolean) => void; + // Actions: models & server setModels: (models: ModelInfo[]) => void; setModelsLoading: (loading: boolean) => void; setSelectedModel: (model: string) => void; setServerInfo: (info: ServerInfo | null) => void; setSavings: (data: SavingsData | null) => void; + incrementSavings: (usage: TokenUsage) => void; + + // Live GPU metrics — streamed from /api/research system_metrics events. + // When non-null, the System panel renders this instead of polled values + // so Power (W) and Energy (kJ) update in real time during a research run. + liveEnergy: LiveEnergyMetrics | null; + setLiveEnergy: (data: LiveEnergyMetrics | null) => void; // Actions: settings updateSettings: (partial: Partial) => void; @@ -387,6 +403,8 @@ export const useAppStore = create((set, get) => { usage?: TokenUsage, telemetry?: MessageTelemetry, audio?: { url: string }, + researchTraces?: ResearchSearchTrace[], + researchSources?: ResearchSource[], ) => { const store = loadConversations(); const conv = store.conversations[conversationId]; @@ -398,6 +416,8 @@ export const useAppStore = create((set, get) => { if (usage) lastMsg.usage = usage; if (telemetry) lastMsg.telemetry = telemetry; if (audio) lastMsg.audio = audio; + if (researchTraces) lastMsg.researchTraces = researchTraces; + if (researchSources) lastMsg.researchSources = researchSources; conv.updatedAt = Date.now(); saveConversations(store); set({ messages: [...conv.messages] }); @@ -412,6 +432,10 @@ export const useAppStore = create((set, get) => { set({ streamState: INITIAL_STREAM }); }, + // ── Deep Research ───────────────────────────────────────────── + deepResearch: false, + setDeepResearch: (on: boolean) => set({ deepResearch: on }), + // ── Models & server ──────────────────────────────────────────── setModels: (models: ModelInfo[]) => set({ models }), @@ -419,6 +443,26 @@ export const useAppStore = create((set, get) => { setSelectedModel: (model: string) => set({ selectedModel: model }), setServerInfo: (info: ServerInfo | null) => set({ serverInfo: info }), setSavings: (data: SavingsData | null) => set({ savings: data }), + incrementSavings: (usage: TokenUsage) => { + const cur = get().savings; + const prompt = usage.prompt_tokens ?? 0; + const completion = usage.completion_tokens ?? 0; + const total = usage.total_tokens ?? prompt + completion; + set({ + savings: { + total_calls: (cur?.total_calls ?? 0) + 1, + total_prompt_tokens: (cur?.total_prompt_tokens ?? 0) + prompt, + total_completion_tokens: (cur?.total_completion_tokens ?? 0) + completion, + total_tokens: (cur?.total_tokens ?? 0) + total, + local_cost: cur?.local_cost ?? 0, + per_provider: cur?.per_provider ?? [], + token_counting_version: cur?.token_counting_version, + }, + }); + }, + + liveEnergy: null, + setLiveEnergy: (data: LiveEnergyMetrics | null) => set({ liveEnergy: data }), cachedConnectors: null, setCachedConnectors: (list) => set({ cachedConnectors: list }), diff --git a/frontend/src/pages/DataSourcesPage.tsx b/frontend/src/pages/DataSourcesPage.tsx index 5de4b908..a5fce03b 100644 --- a/frontend/src/pages/DataSourcesPage.tsx +++ b/frontend/src/pages/DataSourcesPage.tsx @@ -24,7 +24,7 @@ import { import type { LucideIcon } from 'lucide-react'; import { SOURCE_CATALOG } from '../types/connectors'; import type { ConnectRequest } from '../types/connectors'; -import { listConnectors, connectSource, getSyncStatus, triggerSync } from '../lib/connectors-api'; +import { listConnectors, connectSource, disconnectSource, getSyncStatus, triggerSync } from '../lib/connectors-api'; import type { SyncStatus } from '../types/connectors'; // --------------------------------------------------------------------------- @@ -312,11 +312,120 @@ const IconFor = ({ id, size = 18 }: { id: string; size?: number }) => { return ; }; +// The Gmail card unifies the OAuth (`gmail`) and IMAP (`gmail_imap`) backend +// connectors — both should resolve to the gmail_imap catalog entry so the +// connected card shows the same name, unit label, and troubleshooting tips +// regardless of which underlying flow the user picked. +function metaFor(connectorId: string) { + const id = connectorId === 'gmail' ? 'gmail_imap' : connectorId; + return SOURCE_CATALOG.find((s) => s.connector_id === id); +} + +// Advanced OAuth disclosure for the unified Gmail card. Hidden by default; +// expands to a Client ID + Client Secret form that POSTs to the OAuth +// `gmail` backend connector. Lives here rather than in SOURCE_CATALOG +// because the Gmail card is the only one with a dual-flow shape. +function GmailOAuthAdvanced({ + loading, + onConnect, +}: { + loading: boolean; + onConnect: (req: ConnectRequest) => void; +}) { + const [open, setOpen] = useState(false); + return ( + + ); +} + // --------------------------------------------------------------------------- // Data Sources section // --------------------------------------------------------------------------- // Sync status display component with progress bar +function formatTimeAgo(iso: string | null | undefined): string | null { + if (!iso) return null; + const t = new Date(iso).getTime(); + if (Number.isNaN(t)) return null; + const diffSec = (Date.now() - t) / 1000; + if (diffSec < 30) return 'just now'; + if (diffSec < 60) return 'less than a min ago'; + if (diffSec < 3600) { + const m = Math.round(diffSec / 60); + return `${m} min${m === 1 ? '' : 's'} ago`; + } + if (diffSec < 86400) { + const h = Math.round(diffSec / 3600); + return `${h} hr${h === 1 ? '' : 's'} ago`; + } + const d = Math.round(diffSec / 86400); + return `${d} day${d === 1 ? '' : 's'} ago`; +} + +/** Render how far back the corpus extends, given the oldest indexed + * item's timestamp. Returns null when there isn't enough data yet. */ +function formatBacklogRange(iso: string | null | undefined): string | null { + if (!iso) return null; + const t = new Date(iso).getTime(); + if (Number.isNaN(t)) return null; + const days = (Date.now() - t) / 86400_000; + if (days < 7) return 'past few days'; + if (days < 30) return 'past month'; + if (days < 90) return 'past 3 months'; + if (days < 365) return 'past year'; + const years = Math.round(days / 365); + return `past ${years} year${years === 1 ? '' : 's'}`; +} + function SyncStatusDisplay({ chunks, sync, @@ -368,13 +477,65 @@ function SyncStatusDisplay({ ); } - // Done — has chunks - if (chunks > 0) { + // Treat the SyncEngine's checkpointed items_synced as the source of + // truth for "total indexed" — `chunks` from listConnectors counts + // embedding chunks (often != source items) and the checkpoint is what + // both the syncing and idle branches need to display consistently. + const totalIndexed = sync?.items_synced ?? chunks; + const itemsTotal = sync?.items_total ?? 0; + const backlogRange = formatBacklogRange(sync?.oldest_item_date); + // "Complete inbox" — the user has indexed everything reachable. Only + // surface this label when idle (during a sync we always show how far + // back we've gotten so far). + const isComplete = + totalIndexed > 0 && itemsTotal > 0 && totalIndexed >= itemsTotal; + + // Actively syncing — single status line + reassurance line. + if (sync?.state === 'syncing' || syncing) { + const rangeLabel = backlogRange ?? 'building corpus'; + return ( +
+
+ Indexed{' '} + + {totalIndexed.toLocaleString()} {unitLabel} + {' '} + + ({rangeLabel}) + {' '} + + · Still indexing… + +
+
+ Deep Research available now · results improve as more {unitLabel} are indexed +
+
+ ); + } + + // Idle — already has indexed items: show the corpus size + range or + // "complete inbox" label, plus how long ago we last refreshed it. + if (totalIndexed > 0) { + const lastSyncLabel = formatTimeAgo(sync?.last_sync); + const rangeLabel = isComplete + ? 'complete inbox' + : backlogRange; return (
- {chunks.toLocaleString()} {unitLabel} + Indexed {totalIndexed.toLocaleString()} {unitLabel} + {rangeLabel && ( + + {' '}({rangeLabel}) + + )} + {lastSyncLabel && ( + + {' · '}Last synced {lastSyncLabel} + + )}
{connected.map((c) => { - const meta = SOURCE_CATALOG.find(s => s.connector_id === c.connector_id); + const meta = metaFor(c.connector_id); const unit = meta?.unitLabel || 'items'; const sync = syncStatuses[c.connector_id]; - const isReconnecting = expandedId === c.connector_id; const hasError = !!sync?.error; return (
- {c.display_name} + {meta?.display_name ?? c.display_name}
- {isReconnecting && meta?.steps && ( -
-
- Re-enter credentials to reconnect this source. -
- {meta.steps.map((step, i) => ( -
-
- STEP {i + 1} -
-
{step.label}
- {step.url && ( - - {step.urlLabel || 'Open'} → - - )} -
- ))} - {meta.inputFields && ( - handleConnect(c.connector_id, req)} - /> - )} -
- )}
); })} @@ -744,7 +850,7 @@ function DataSourcesSection() {
{notConnected.map((c) => { - const meta = SOURCE_CATALOG.find(s => s.connector_id === c.connector_id); + const meta = metaFor(c.connector_id); const isExpanded = expandedId === c.connector_id; return ( @@ -767,10 +873,10 @@ function DataSourcesSection() { >
- {c.display_name} + {meta?.display_name ?? c.display_name}
- Not connected + {meta?.description ?? 'Not connected'}
@@ -822,6 +928,12 @@ function DataSourcesSection() { onSubmit={(req) => handleConnect(c.connector_id, req)} /> )} + {c.connector_id === 'gmail_imap' && ( + handleConnect('gmail', req)} + /> + )} {meta?.troubleshooting && (
diff --git a/frontend/src/types/connectors.ts b/frontend/src/types/connectors.ts index cb212be6..534a2a80 100644 --- a/frontend/src/types/connectors.ts +++ b/frontend/src/types/connectors.ts @@ -36,6 +36,13 @@ export interface SyncStatus { state: "idle" | "syncing" | "paused" | "error"; items_synced: number; items_total: number; + /** Items processed in the current (or most recent) run only. `null` + * when no sync has been triggered through this server session yet. */ + new_items_synced?: number | null; + /** ISO 8601 timestamp of the oldest indexed item, used to label how far + * back the corpus reaches ("past 3 months", "past 5 years"). `null` + * before anything is indexed. */ + oldest_item_date?: string | null; last_sync: string | null; error: string | null; } @@ -73,6 +80,9 @@ export const SOURCE_CATALOG: ConnectorMeta[] = [ }, // ── Communication ────────────────────────────────────────────────── { + // Unified Gmail card. Defaults to the IMAP (app-password) flow because + // it needs no Google Cloud setup; the OAuth path is offered as an + // "Advanced" disclosure rendered in DataSourcesPage. connector_id: 'gmail_imap', display_name: 'Gmail', auth_type: 'oauth', @@ -83,23 +93,14 @@ export const SOURCE_CATALOG: ConnectorMeta[] = [ unitLabel: 'emails', steps: [ { - label: 'Go to your Google Account \u2192 Security \u2192 2-Step Verification. Make sure it\'s turned ON. App Passwords only work if 2-Step Verification is enabled.', - url: 'https://myaccount.google.com/signinoptions/two-step-verification', - urlLabel: 'Open Google Security \u2192', - }, - { - label: 'Go to App Passwords. Select app: "Mail", device: "Other" and type "OpenJarvis". Click Generate. This does NOT open a login popup \u2014 you\'ll get a 16-character password to copy (you won\'t see it again).', + label: 'Make sure 2-Step Verification is enabled, then generate a 16-character App Password (Mail / Other / "OpenJarvis"). Paste it below \u2014 spaces are fine, and use the app password, not your regular Gmail password.', url: 'https://myaccount.google.com/apppasswords', - urlLabel: 'Open App Passwords \u2192', - }, - { - label: 'Paste your Gmail address and the 16-character app password below (spaces are fine). Use the app password, NOT your regular Gmail password.', + urlLabel: 'How to get an app password \u2192', }, ], troubleshooting: [ "Don't see App Passwords? Make sure 2-Step Verification is enabled first.", "Google Workspace user? Your admin may need to enable App Passwords for your organization.", - "Want OAuth instead? Use the Google Drive connector \u2014 it covers Gmail content too.", ], inputFields: [ { name: 'email', placeholder: 'you@gmail.com', type: 'text' }, diff --git a/frontend/src/types/index.ts b/frontend/src/types/index.ts index 923fe165..4ecbb62d 100644 --- a/frontend/src/types/index.ts +++ b/frontend/src/types/index.ts @@ -61,12 +61,69 @@ export interface MessageTelemetry { suggested_max_tokens?: number; } +export interface TimeRange { + start?: string; + end?: string; +} + +export interface ResearchSource { + ref: number; + title?: string; + sender?: string; + date?: string; + url?: string; +} + +export interface ResearchSearchTrace { + id: string; + query: string; + person?: string; + timeRange?: TimeRange | string; + status: 'pending' | 'complete'; + numHits?: number; + topTitles?: string[]; +} + +export type ResearchEvent = + | { + type: 'search_call'; + arguments: { + query: string; + person?: string; + time_range?: TimeRange | string; + }; + } + | { + type: 'search_result'; + num_hits: number; + top_titles?: string[]; + sources?: ResearchSource[]; + } + | { type: 'synthesis'; text: string } + | { + type: 'system_metrics'; + power_w: number; + energy_j: number; + duration_s: number; + } + | { type: 'done'; usage?: TokenUsage } + | { type: 'error'; message: string }; + +export interface LiveEnergyMetrics { + power_w: number; + energy_j: number; + duration_s: number; +} + export interface ChatMessage { id: string; role: 'user' | 'assistant'; content: string; timestamp: number; toolCalls?: ToolCallInfo[]; + researchTraces?: ResearchSearchTrace[]; + researchSources?: ResearchSource[]; + isResearch?: boolean; usage?: TokenUsage; telemetry?: MessageTelemetry; audio?: { url: string }; diff --git a/frontend/vite.config.ts b/frontend/vite.config.ts index cccf520f..2cc44afc 100644 --- a/frontend/vite.config.ts +++ b/frontend/vite.config.ts @@ -29,7 +29,7 @@ export default defineConfig({ }, workbox: { globPatterns: ['**/*.{js,css,html,ico,png,svg}'], - navigateFallbackDenylist: [/^\/v1\//, /^\/health/, /^\/dashboard/], + navigateFallbackDenylist: [/^\/v1\//, /^\/health/, /^\/dashboard/, /^\/api\//], }, }), ], @@ -53,6 +53,7 @@ export default defineConfig({ proxy: { '/v1': process.env.VITE_API_URL || 'http://localhost:8000', '/health': process.env.VITE_API_URL || 'http://localhost:8000', + '/api': process.env.VITE_API_URL || 'http://localhost:8000', }, }, }); diff --git a/pyproject.toml b/pyproject.toml index 159f6edf..1bc774b8 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -30,6 +30,7 @@ dependencies = [ "httpx>=0.27", "openai>=1.30", "posthog>=3.0", + "pynvml>=13.0.1", "python-telegram-bot>=22.6", "rich>=13", "tomli>=2.0; python_version < '3.11'", @@ -183,6 +184,10 @@ select = ["E", "F", "I", "W"] # hybrid/ is research code with long prompt strings and paradigm-specific # config dicts — same relaxation as evals research code above. "src/openjarvis/agents/hybrid/*.py" = ["E501"] +# research_loop.py carries the multi-paragraph planner system prompt as +# inline string literals; line-length wrapping would harm readability of +# the prompt itself. +"src/openjarvis/agents/research_loop.py" = ["E501"] [dependency-groups] dev = [ diff --git a/src/openjarvis/agents/executor.py b/src/openjarvis/agents/executor.py index 6a5fe6e0..8d6be76f 100644 --- a/src/openjarvis/agents/executor.py +++ b/src/openjarvis/agents/executor.py @@ -93,20 +93,34 @@ class AgentExecutor: ) return agent.run(input_text) - def execute_tick(self, agent_id: str) -> None: + def execute_tick( + self, agent_id: str, *, lock_already_held: bool = False + ) -> None: """Run one tick for the given agent. 1. Acquire concurrency guard (start_tick) 2. Invoke agent with retry logic 3. Update stats 4. Release guard (end_tick) + + ``lock_already_held`` is set by callers that took the start_tick() + lock themselves before spawning the worker (e.g. the HTTP /run route + guards against concurrent POSTs by acquiring before threading). + Without this flag, the executor would re-acquire and trip its own + guard — bailing out with no end_tick(), leaving the agent stuck in + ``status='running'`` forever. """ - try: - self._manager.start_tick(agent_id) + if lock_already_held: self._set_activity(agent_id, "Preparing tick...") - except ValueError: - logger.warning("Agent %s already running, skipping tick", agent_id) - return + else: + try: + self._manager.start_tick(agent_id) + self._set_activity(agent_id, "Preparing tick...") + except ValueError: + logger.warning( + "Agent %s already running, skipping tick", agent_id + ) + return agent = self._manager.get_agent(agent_id) if agent is None: diff --git a/src/openjarvis/agents/manager.py b/src/openjarvis/agents/manager.py index 5b03d38d..6eed3302 100644 --- a/src/openjarvis/agents/manager.py +++ b/src/openjarvis/agents/manager.py @@ -7,6 +7,7 @@ to the five existing primitives (Intelligence, Agent, Tools, Engine, Learning). from __future__ import annotations import json +import logging import sqlite3 import time import uuid @@ -14,6 +15,8 @@ from pathlib import Path from typing import Any, Dict, List, Optional from uuid import uuid4 +logger = logging.getLogger(__name__) + _CREATE_AGENTS = """\ CREATE TABLE IF NOT EXISTS managed_agents ( id TEXT PRIMARY KEY, @@ -123,6 +126,33 @@ class AgentManager: except sqlite3.OperationalError: pass # Column already exists self._conn.commit() + self._clear_stale_running_state() + + def _clear_stale_running_state(self) -> None: + """Reset any agent stuck in ``status='running'`` on startup. + + Tick worker threads are ``daemon=True`` — when the server process + exits (SIGTERM, crash, restart), they die without running the + ``finally`` clause that calls :meth:`end_tick`, leaving the DB row + in ``running`` forever. The :meth:`start_tick` guard then rejects + every subsequent run with "Agent is already running". + + A freshly-started process holds zero tick locks by definition, so + any persisted ``running`` is a zombie. Sweep it back to ``idle`` + and clear the activity string so the UI doesn't show a stale + "Preparing tick..." indicator. + """ + cur = self._conn.execute( + "UPDATE managed_agents SET status = 'idle', current_activity = ''," + " updated_at = ? WHERE status = 'running'", + (time.time(),), + ) + self._conn.commit() + if cur.rowcount: + logger.info( + "AgentManager: cleared stale 'running' status on %d agent(s)", + cur.rowcount, + ) def close(self) -> None: self._conn.close() diff --git a/src/openjarvis/agents/research_loop.py b/src/openjarvis/agents/research_loop.py new file mode 100644 index 00000000..3a1e85ae --- /dev/null +++ b/src/openjarvis/agents/research_loop.py @@ -0,0 +1,683 @@ +"""Agentic research loop over the hybrid-search tool. + +A small, self-contained planner-executor loop: + +* the planner is a local Ollama chat model (default ``gemma4:31b``), +* the only tool it can call is :meth:`HybridSearch.search`, +* it gets up to ``max_iterations`` tool calls, +* tool results are trimmed before re-entering the context window, and +* the final reply must cite specific hits. + +The loop is deliberately decoupled from the rest of the agent scaffolding +(`ToolUsingAgent`, `EventBus`, `AgentContext`, etc.) so the surface stays +small. Anything that wants tracing or registry integration can wrap it. +""" + +from __future__ import annotations + +import json +import logging +import sys +from dataclasses import dataclass, field +from datetime import datetime +from typing import Any, Callable, Dict, List, Optional + +from openjarvis.connectors.hybrid_search import HybridSearch, SearchHit +from openjarvis.core.types import Message, Role, ToolCall +from openjarvis.engine._base import InferenceEngine + +logger = logging.getLogger(__name__) + + +DEFAULT_PLANNER_MODEL = "gemma4:31b" + + +CLARIFY_TOOL_SPEC: Dict[str, Any] = { + "type": "function", + "function": { + "name": "clarify", + "description": ( + "Ask the user a clarifying question and wait for their answer. " + "Only use AFTER at least one search has been attempted. Use when " + "search results are ambiguous (e.g. three different people share " + "a first name), search returned zero results and the query likely " + "needs reframing, or the scope is too broad to synthesize " + "meaningfully. Never use clarify before searching." + ), + "parameters": { + "type": "object", + "properties": { + "question": { + "type": "string", + "description": ( + "The clarifying question to ask the user. Be specific " + "about what you need to know to make progress." + ), + }, + }, + "required": ["question"], + }, + }, +} + + +SEARCH_TOOL_SPEC: Dict[str, Any] = { + "type": "function", + "function": { + "name": "search", + "description": ( + "Hybrid search over the user's personal knowledge corpus (emails, " + "notes, calendar events, attachments). Combines BM25 lexical match " + "with dense embedding similarity, ranked by reciprocal rank fusion. " + "Use structured filters (person, time_range, sources) whenever the " + "user names a specific person or time window. Each call returns up " + "to 'limit' results with content snippets and thread context." + ), + "parameters": { + "type": "object", + "properties": { + "query": { + "type": "string", + "description": ( + "Natural-language query. Use the topic the user is asking " + "about. Can be empty when filtering purely by person or " + "time (e.g. 'list all mail from Kelly in May')." + ), + }, + "person": { + "type": "string", + "description": ( + "Filter to messages involving this person. Matches a " + "substring of the name or email address — 'Kelly' or " + "'@tldrnewsletter.com' both work." + ), + }, + "time_range": { + "type": "object", + "description": "ISO 8601 datetime range. Either bound may be omitted.", + "properties": { + "start": {"type": "string", "description": "ISO 8601 start"}, + "end": {"type": "string", "description": "ISO 8601 end"}, + }, + }, + "sources": { + "type": "array", + "description": "Restrict to these connectors (e.g. ['gmail']).", + "items": {"type": "string"}, + }, + "limit": { + "type": "integer", + "description": "Max results to return (default 20, cap 20).", + "default": 20, + }, + }, + "required": ["query"], + }, + }, +} + + +SYSTEM_PROMPT = """You are a research assistant with access to the user's personal knowledge corpus (their email, notes, calendar). You answer questions by calling two tools: + + search(query, person=None, time_range=None, sources=None, limit=20) + clarify(question) + +Strategy: + 1. If the user names a person, ALWAYS pass `person=` rather than relying on lexical match. Hybrid search will fuzzy-match name or address fragments. + 2. When the user mentions ANY time window — "this past week", "recently", "last month", "past few days", "yesterday" — you MUST translate it to a `time_range` parameter. Today is {today}. + 3. The `time_range` argument is a JSON object: `{{"start": "", "end": ""}}`. Either bound may be omitted, but pass at least one whenever the user gave you a temporal cue. + 4. If the first structured search returns nothing useful, broaden with a semantic query and drop filters one at a time. + 5. You have a clarify tool. Only use it AFTER at least one search attempt. Use it when: you found multiple ambiguous matches (e.g. 3 different people named John), search returned zero results and the query might need reframing, or the scope is too broad to synthesize meaningfully. Never use clarify before searching — always try first. + 6. After receiving a clarify response, use the information to construct a precise search with the correct person, time_range, and query parameters. Never send an empty query or a query with no parameters — extract every concrete signal from the user's reply (names, dates, topics) and put it on the call. + 7. Tool calls — search AND clarify — share a budget of 5 total. Spend wisely. + +Synthesis rules: + - Cite sources as individual numbers in square brackets. Always separate — write [4] [7] [20], never [4, 7, 20]. Never format citations as markdown links. Just the number in brackets: [1]. The `ref` field on each hit is the citation number. + - Quote sender / date / subject when relevant — the user wants attribution. + - If the search returned nothing relevant, say so plainly. Do not invent results. + - Only state facts that appear in the retrieved search results. Never supplement with your own knowledge or training data. If you are unsure whether a fact came from the search results, do not include it. + +Today's date is {today}. +""" + + +# --------------------------------------------------------------------------- +# Tool-result shaping +# --------------------------------------------------------------------------- + + +def _trim_thread_context(ctx: List[Dict[str, Any]], cap: int) -> List[Dict[str, Any]]: + """Keep the first ``cap`` entries; mark elision when trimming.""" + if len(ctx) <= cap: + return ctx + trimmed = list(ctx[:cap]) + trimmed.append({"snippet": f"… {len(ctx) - cap} more chunks in thread …"}) + return trimmed + + +def shape_results_for_model( + hits: List[SearchHit], + *, + detailed_top: int = 5, + thread_ctx_per_hit: int = 3, + total_cap: int = 20, +) -> Dict[str, Any]: + """Compact a hit list into a JSON payload the planner can chew through. + + The first ``detailed_top`` rows keep their content snippet and trimmed + thread context; the remainder are summarised to title + sender + date so + the planner still sees the breadth of what's available without blowing the + context window. Each hit gets a 1-indexed numeric ``ref`` so the synthesis + can cite it as ``[N]``. + """ + out_hits: List[Dict[str, Any]] = [] + visible = hits[:total_cap] + for i, h in enumerate(visible): + sender = h.participants[0] if h.participants else "" + base = { + "ref": i + 1, + "title": h.title, + "sender": sender, + "timestamp": h.timestamp, + "source": h.source, + "score": round(h.score, 4), + } + if i < detailed_top: + base["snippet"] = h.content_snippet + if h.thread_context: + base["thread"] = _trim_thread_context(h.thread_context, thread_ctx_per_hit) + out_hits.append(base) + return { + "num_results": len(hits), + "shown": len(visible), + "truncated": len(hits) > total_cap, + "hits": out_hits, + } + + +def _hit_date(timestamp: str) -> str: + """Pull a ``YYYY-MM-DD`` date out of a SearchHit timestamp (best effort).""" + if not timestamp: + return "" + try: + return datetime.fromisoformat(timestamp.replace("Z", "+00:00")).date().isoformat() + except (ValueError, AttributeError): + return str(timestamp)[:10] + + +def _bare_doc_id(source: str, document_id: str) -> str: + """Strip the connector prefix from a stored ``doc_id``. + + Gmail ingest writes ``doc_id="gmail:"`` so that ids stay + unique across connectors. The Gmail web UI only resolves the bare hex + message id; passing the full prefixed form 404s and bounces the user + back to the inbox. Other connectors can be added here as we link out + to them. + """ + if not document_id: + return "" + prefix = f"{source}:" + if source and document_id.startswith(prefix): + return document_id[len(prefix):] + return document_id + + +def _hit_url(source: str, document_id: str) -> str: + """Build a clickable URL for a hit, when we know how to link it. + + Gmail is the only source we currently link out for; everything else + returns an empty string and the client falls back to non-clickable + citation chips. + + Two id flavors land here: + + - **Hex message id** (``19dfa2ccbeff78b0``) — what the OAuth Gmail + connector stores. Resolves directly via ``#all/`` permalink. + - **RFC822 Message-ID** (````) — what the IMAP + connector stores, since IMAP doesn't expose Gmail's internal hex + id. The permalink form would 404; instead route through Gmail's + search URL with the ``rfc822msgid:`` operator, which lands the user + on the specific message. + """ + if source == "gmail" and document_id: + msg_id = _bare_doc_id(source, document_id) + if not msg_id: + return "" + if "@" in msg_id or "<" in msg_id or ">" in msg_id: + rfc_id = msg_id.strip("<>") + return f"https://mail.google.com/mail/u/0/#search/rfc822msgid:{rfc_id}" + return f"https://mail.google.com/mail/u/0/#all/{msg_id}" + return "" + + +def build_sources_for_client( + hits: List[SearchHit], + *, + total_cap: int = 20, +) -> List[Dict[str, Any]]: + """Produce the citation-friendly sources list streamed to the frontend. + + One entry per hit, in the same order the planner sees them — so a + ``[N]`` citation in the synthesis maps to ``sources[N - 1]`` on the + client. We don't deduplicate by ``document_id``: separate chunks of the + same email each get their own citation slot since the planner may quote + different parts. + """ + out: List[Dict[str, Any]] = [] + for i, h in enumerate(hits[:total_cap]): + sender = h.participants[0] if h.participants else "" + out.append( + { + "ref": i + 1, + "title": h.title, + "sender": sender, + "date": _hit_date(h.timestamp), + "source_id": _bare_doc_id(h.source, h.document_id), + "url": _hit_url(h.source, h.document_id), + } + ) + return out + + +# --------------------------------------------------------------------------- +# Agent +# --------------------------------------------------------------------------- + + +@dataclass +class ToolInvocation: + """One tool call together with what the planner asked for and got. + + ``tool_name`` is ``"search"`` or ``"clarify"``. For search calls, + ``num_results``, ``top_titles`` and ``raw_hits`` are populated; for + clarify calls, ``response`` holds the user's answer. + """ + + arguments: Dict[str, Any] + num_results: int = 0 + top_titles: List[str] = field(default_factory=list) + raw_hits: List[SearchHit] = field(default_factory=list) + tool_name: str = "search" + response: str = "" + + +def _default_clarify_handler(question: str) -> str: + """Prompt the user on stdout and read a one-line answer from stdin. + + Empty answers are echoed back as a sentinel so the planner doesn't think + the user was silent because of an upstream error. + """ + print(file=sys.stderr) + print(f"\033[1m🤔 Clarification needed:\033[0m {question}", file=sys.stderr) + try: + answer = input("> ").strip() + except EOFError: + return "(no answer provided)" + return answer or "(user did not provide a clarification)" + + +@dataclass +class ResearchResult: + answer: str + iterations: int + tool_calls: List[ToolInvocation] + usage: Dict[str, int] = field(default_factory=dict) + + +class ResearchAgent: + """Planner + executor loop over a single hybrid-search tool. + + Parameters + ---------- + engine: + An ``InferenceEngine`` that supports OpenAI-style ``tools`` in + ``generate`` (Ollama with a tool-capable model). + search: + The HybridSearch instance the planner can call. + model: + Planner model tag (default ``gemma4:31b``). + max_iterations: + Hard ceiling on tool calls before the loop is forced into synthesis. + temperature, max_tokens, num_ctx: + Generation parameters passed through to ``engine.generate``. + on_event: + Optional callback fired at loop milestones so callers (e.g. the SSE + research router) can stream progress without rewriting the loop. + Receives a dict in one of these shapes: + - ``{"type": "search_call", "arguments": {...}}`` — about to call search + - ``{"type": "search_result", "num_hits": N, "top_titles": [...], "sources": [{"ref": 1, "title": ..., "sender": ..., "date": ..., "source_id": ..., "url": ...}, ...]}`` — search returned + - ``{"type": "clarify_call", "question": "..."}`` — about to ask for clarification + - ``{"type": "clarify_response", "response": "..."}`` — clarification received + - ``{"type": "final_answer", "text": "..."}`` — synthesis ready + The callback runs on the same thread as ``run`` and must be non-blocking. + """ + + def __init__( + self, + engine: InferenceEngine, + search: HybridSearch, + *, + model: str = DEFAULT_PLANNER_MODEL, + max_iterations: int = 5, + temperature: float = 0.3, + max_tokens: int = 1500, + num_ctx: int = 16384, + clarify_handler: Optional[Callable[[str], str]] = None, + on_event: Optional[Callable[[Dict[str, Any]], None]] = None, + ) -> None: + self._engine = engine + self._search = search + self._model = model + self._max_iterations = int(max_iterations) + self._temperature = float(temperature) + self._max_tokens = int(max_tokens) + self._num_ctx = int(num_ctx) + self._clarify_handler = clarify_handler or _default_clarify_handler + self._on_event = on_event + + def _emit(self, event: Dict[str, Any]) -> None: + """Fire ``self._on_event`` if set; swallow callback errors.""" + if self._on_event is None: + return + try: + self._on_event(event) + except Exception as exc: # noqa: BLE001 + logger.debug("on_event callback raised %s — ignoring", exc) + + # ------------------------------------------------------------------ + # Argument parsing + # ------------------------------------------------------------------ + + @staticmethod + def _parse_time_range(raw: Any): + if not raw or not isinstance(raw, dict): + return None + def _maybe(v): + if not v: + return None + try: + return datetime.fromisoformat(str(v).replace("Z", "+00:00")) + except ValueError: + return None + start = _maybe(raw.get("start")) + end = _maybe(raw.get("end")) + if start is None and end is None: + return None + return (start, end) + + def _execute_search(self, args: Dict[str, Any]) -> ToolInvocation: + query = str(args.get("query", "") or "") + person = args.get("person") or None + time_range = self._parse_time_range(args.get("time_range")) + sources = args.get("sources") or None + if sources and not isinstance(sources, list): + sources = [str(sources)] + limit = int(args.get("limit", 20) or 20) + limit = max(1, min(limit, 20)) + + hits = self._search.search( + query, + person=person, + time_range=time_range, + sources=sources, + limit=limit, + ) + titles = [h.title or (h.content_snippet[:60] + "…") for h in hits[:5]] + return ToolInvocation( + tool_name="search", + arguments={ + "query": query, + "person": person, + "time_range": ( + {"start": time_range[0].isoformat() if time_range and time_range[0] else None, + "end": time_range[1].isoformat() if time_range and time_range[1] else None} + if time_range else None + ), + "sources": sources, + "limit": limit, + }, + num_results=len(hits), + top_titles=titles, + raw_hits=hits, + ) + + def _execute_clarify(self, args: Dict[str, Any]) -> ToolInvocation: + question = str(args.get("question", "") or "").strip() + if not question: + return ToolInvocation( + tool_name="clarify", + arguments={"question": ""}, + response="(no question provided by agent — skipping clarify)", + ) + answer = self._clarify_handler(question) + return ToolInvocation( + tool_name="clarify", + arguments={"question": question}, + response=answer, + ) + + # ------------------------------------------------------------------ + # Loop + # ------------------------------------------------------------------ + + def run(self, query: str) -> ResearchResult: + """Run the loop end-to-end and return the synthesis plus a trace.""" + sys_msg = Message( + role=Role.SYSTEM, + content=SYSTEM_PROMPT.format(today=datetime.now().isoformat(timespec="minutes")), + ) + messages: List[Message] = [sys_msg, Message(role=Role.USER, content=query)] + + invocations: List[ToolInvocation] = [] + total_usage = {"prompt_tokens": 0, "completion_tokens": 0, "total_tokens": 0} + + iterations = 0 + for _ in range(self._max_iterations + 1): + iterations += 1 + tools_arg = ( + [SEARCH_TOOL_SPEC, CLARIFY_TOOL_SPEC] + if len(invocations) < self._max_iterations + else None + ) + result = self._engine.generate( + messages, + model=self._model, + temperature=self._temperature, + max_tokens=self._max_tokens, + num_ctx=self._num_ctx, + tools=tools_arg, + ) + for k in total_usage: + total_usage[k] += int(result.get("usage", {}).get(k, 0)) + + content = result.get("content", "") or "" + tool_calls_raw = result.get("tool_calls", []) or [] + + if not tool_calls_raw: + if content.strip(): + answer = content.strip() + self._emit({"type": "final_answer", "text": answer}) + return ResearchResult( + answer=answer, + iterations=iterations, + tool_calls=invocations, + usage=total_usage, + ) + # Empty content with no tool call — push a synthesis prod + if invocations: + messages.append(Message(role=Role.ASSISTANT, content=content)) + messages.append( + Message( + role=Role.USER, + content=( + "Write your final answer now based on the search " + "results above. Cite sources as [1], [2], etc." + ), + ) + ) + continue + fallback = "(model returned no content and no tool calls)" + self._emit({"type": "final_answer", "text": fallback}) + return ResearchResult( + answer=fallback, + iterations=iterations, + tool_calls=invocations, + usage=total_usage, + ) + + assistant_msg = Message( + role=Role.ASSISTANT, + content=content, + tool_calls=[ + ToolCall( + id=tc.get("id", f"call_{i}"), + name=tc.get("name", "search"), + arguments=tc.get("arguments", "{}") or "{}", + ) + for i, tc in enumerate(tool_calls_raw) + ], + ) + messages.append(assistant_msg) + + for tc in tool_calls_raw: + name = tc.get("name", "") + raw_args = tc.get("arguments", "{}") or "{}" + try: + args = json.loads(raw_args) if isinstance(raw_args, str) else dict(raw_args) + except json.JSONDecodeError: + args = {} + + if name == "search": + # Guard against the planner pre-empting clarify before any + # search has run — silently accept; the rule lives in the + # system prompt as guidance, not enforcement. + self._emit({"type": "search_call", "arguments": args}) + inv = self._execute_search(args) + invocations.append(inv) + self._emit( + { + "type": "search_result", + "num_hits": inv.num_results, + "top_titles": inv.top_titles, + "sources": build_sources_for_client(inv.raw_hits), + } + ) + tool_output = json.dumps( + shape_results_for_model(inv.raw_hits), ensure_ascii=False + ) + elif name == "clarify": + # Enforce the "search first" rule at runtime so we don't + # surprise the user with a clarification before showing any + # work. If the planner jumps to clarify with no searches + # behind it, return an error and let the loop try again. + if not any(i.tool_name == "search" for i in invocations): + tool_output = json.dumps( + { + "error": ( + "clarify is only available after at least " + "one search call. Run search first, then " + "use clarify if the results are ambiguous " + "or empty." + ) + } + ) + else: + self._emit( + {"type": "clarify_call", "question": str(args.get("question", ""))} + ) + inv = self._execute_clarify(args) + invocations.append(inv) + self._emit( + {"type": "clarify_response", "response": inv.response} + ) + tool_output = json.dumps( + { + "question": inv.arguments.get("question", ""), + "user_response": inv.response, + } + ) + else: + tool_output = json.dumps( + { + "error": ( + f"unknown tool {name!r}; available tools are " + "'search' and 'clarify'" + ) + } + ) + + messages.append( + Message( + role=Role.TOOL, + content=tool_output, + tool_call_id=tc.get("id", ""), + name=name, + ) + ) + + if len(invocations) >= self._max_iterations: + messages.append( + Message( + role=Role.USER, + content=( + "You have used your tool-call budget (search + " + "clarify combined). Write the final synthesis now " + "using only the search results and clarifications " + "above. Cite sources as [1], [2], etc." + ), + ) + ) + + # Loop fell through without the model producing a text response. + # Force one final tool-less synthesis call so the caller always gets + # an answer — bailing out with a sentinel string is never useful to + # the user, who already paid for the searches. + messages.append( + Message( + role=Role.USER, + content=( + "You've used all your search attempts. Synthesize your " + "findings now from whatever you've found so far. Do not " + "request more tool calls — write the final answer as " + "plain text, citing sources as [1], [2], etc. where you can. " + "If the searches returned nothing usable, say so plainly." + ), + ) + ) + iterations += 1 + final = self._engine.generate( + messages, + model=self._model, + temperature=self._temperature, + max_tokens=self._max_tokens, + num_ctx=self._num_ctx, + tools=None, + ) + for k in total_usage: + total_usage[k] += int(final.get("usage", {}).get(k, 0)) + answer = (final.get("content", "") or "").strip() + if not answer: + answer = ( + "(no synthesis available — the search budget was exhausted " + "and the model returned no text response)" + ) + self._emit({"type": "final_answer", "text": answer}) + return ResearchResult( + answer=answer, + iterations=iterations, + tool_calls=invocations, + usage=total_usage, + ) + + +__all__ = [ + "ResearchAgent", + "ResearchResult", + "ToolInvocation", + "SEARCH_TOOL_SPEC", + "CLARIFY_TOOL_SPEC", + "SYSTEM_PROMPT", + "DEFAULT_PLANNER_MODEL", + "shape_results_for_model", + "build_sources_for_client", +] diff --git a/src/openjarvis/analytics/identity.py b/src/openjarvis/analytics/identity.py index 63f4453f..3a3f6a3b 100644 --- a/src/openjarvis/analytics/identity.py +++ b/src/openjarvis/analytics/identity.py @@ -9,6 +9,8 @@ No email, no name, no hardware fingerprint — just an opaque UUID. from __future__ import annotations +import os +import sys import uuid from pathlib import Path @@ -45,5 +47,16 @@ def reset_anon_id(path: Path | str) -> str: def is_analytics_enabled(cfg: AnalyticsConfig) -> bool: - """Return True if analytics is enabled in config.""" + """Return True if analytics is enabled in config. + + Always disabled when running under pytest. The PostHog SDK registers + an ``atexit`` hook that synchronously joins its consumer thread; if + the host is unreachable (CI runners can't reach the analytics endpoint), + each queued batch retries for ``timeout * max_retries`` seconds and the + interpreter never exits. Detect pytest via ``PYTEST_CURRENT_TEST`` + (set per test) and ``"pytest" in sys.modules`` (covers the collection + phase before the first test runs). + """ + if os.environ.get("PYTEST_CURRENT_TEST") or "pytest" in sys.modules: + return False return cfg.enabled diff --git a/src/openjarvis/cli/__init__.py b/src/openjarvis/cli/__init__.py index 658682f4..50a85fdc 100644 --- a/src/openjarvis/cli/__init__.py +++ b/src/openjarvis/cli/__init__.py @@ -60,8 +60,14 @@ def cli(ctx: click.Context, verbose: bool, quiet: bool) -> None: ctx.obj["quiet"] = quiet setup_logging(verbose=verbose, quiet=quiet) - # Check for updates on interactive commands - if not quiet and ctx.invoked_subcommand: + # Check for updates on interactive commands. The banner is noise in + # demo recordings of ``jarvis ask --research``, so skip it whenever + # the research flag is in argv (cheap argv sniff — Click hasn't + # parsed the subcommand's args yet at this point). + import sys + + research_mode_active = "--research" in sys.argv + if not quiet and ctx.invoked_subcommand and not research_mode_active: from openjarvis.cli._version_check import check_for_updates check_for_updates(ctx.invoked_subcommand) diff --git a/src/openjarvis/cli/_version_check.py b/src/openjarvis/cli/_version_check.py index d95f4933..a1db0fa4 100644 --- a/src/openjarvis/cli/_version_check.py +++ b/src/openjarvis/cli/_version_check.py @@ -4,6 +4,7 @@ from __future__ import annotations import json import logging +import os import sys import time from pathlib import Path @@ -17,7 +18,13 @@ _CHECK_COMMANDS = {"ask", "chat", "serve"} def check_for_updates(command_name: str) -> None: - """Print a message if a newer version is available. Best-effort, never raises.""" + """Print a message if a newer version is available. Best-effort, never raises. + + Honors ``OPENJARVIS_NO_UPDATE_CHECK`` — set it to any non-empty value + (``1``, ``true``, etc.) to skip the GitHub poll and the banner. + """ + if os.environ.get("OPENJARVIS_NO_UPDATE_CHECK", "").strip(): + return if command_name not in _CHECK_COMMANDS: return try: diff --git a/src/openjarvis/cli/ask.py b/src/openjarvis/cli/ask.py index 8c56499d..b1460ae1 100644 --- a/src/openjarvis/cli/ask.py +++ b/src/openjarvis/cli/ask.py @@ -32,6 +32,187 @@ from openjarvis.telemetry.store import TelemetryStore logger = logging.getLogger(__name__) +def _run_research( + *, + query_text: str, + engine, + model_name: str | None, + knowledge_db: str | None, + output_json: bool, + console: Console, +) -> None: + """Run the hybrid-search research loop and print the result to the console. + + Lazy imports keep the cost of this branch off the cold-path of plain + ``jarvis ask`` calls. + """ + import re + + from rich.markdown import Markdown + from rich.theme import Theme + + from openjarvis.agents.research_loop import DEFAULT_PLANNER_MODEL, ResearchAgent + from openjarvis.connectors.embeddings import OllamaEmbedder + from openjarvis.connectors.hybrid_search import HybridSearch + from openjarvis.connectors.store import KnowledgeStore + from openjarvis.engine.ollama import OllamaEngine + + store_kwargs: dict = {} + if knowledge_db: + store_kwargs["db_path"] = knowledge_db + store = KnowledgeStore(**store_kwargs) + + # Research mode is wired specifically to Ollama: the planner prompt + # (gemma4:31b) and the function-call schema for search/clarify both + # assume Ollama's /api/chat tool semantics. Using the engine returned + # by get_engine() here is a foot-gun — discovery can pick any + # OpenAI-compatible engine registered on the same port as our own + # API server. research_router.py hardcodes OllamaEngine() for the + # same reason; mirror that here so CLI and HTTP behave identically. + engine = OllamaEngine() + + chunk_count = store._conn.execute( + "SELECT COUNT(*) FROM knowledge_chunks" + ).fetchone()[0] + logger.debug( + "research: engine=%s.%s db=%s chunks=%d", + type(engine).__module__, + type(engine).__name__, + store._db_path, + chunk_count, + ) + + embedder = OllamaEmbedder() + embedder_available = embedder.is_available() + logger.debug("research: embedder available=%s", embedder_available) + if not embedder_available: + console.print( + "[yellow]Ollama embedder unavailable — falling back to BM25-only " + "retrieval. Run `ollama pull nomic-embed-text` for hybrid scoring.[/yellow]" + ) + embedder = None + + planner_model = model_name or DEFAULT_PLANNER_MODEL + logger.debug("research: planner_model=%s", planner_model) + + # ---- Output styling -------------------------------------------------- + # Two consoles by design: traces and the timing footer go to stderr + # (so ``jarvis ask --research "..." > out.md`` still gives a clean + # markdown file), while the rendered synthesis goes to stdout. The + # ``markdown.code`` theme override is the cyan-citation hack — see + # ``_style_citations`` below. + trace = Console(stderr=True, soft_wrap=True, highlight=False) + answer_console = Console( + theme=Theme({"markdown.code": "cyan bold not italic"}), + soft_wrap=True, + highlight=False, + ) + + def _style_citations(text: str) -> str: + """Wrap each ``[N]`` in inline code so it renders cyan. + + Rich's Markdown class doesn't expose any hook for styling + arbitrary text spans, so we cheat: convert the citation tokens + into inline-code markdown (`[1]` → `` `[1]` ``) and override + the ``markdown.code`` theme entry above to colour them. Trims + the implicit monospace background that some terminal themes + give inline code so the result reads as text, not as code. + """ + return re.sub(r"(\[\d+\])", r"`\1`", text) + + def _format_search_call(args: dict) -> str: + q = args.get("query", "") or "" + extras: list[str] = [] + person = args.get("person") + if person: + extras.append(f"person: {person}") + time_range = args.get("time_range") + if isinstance(time_range, dict): + start = time_range.get("start") or "" + end = time_range.get("end") or "" + if start or end: + bounds = f"{start or '…'} → {end or '…'}" + extras.append(f"when: {bounds}") + suffix = f" ({', '.join(extras)})" if extras else "" + return f"'{q}'{suffix}" + + def on_event(event: dict) -> None: + etype = event.get("type") + if etype == "search_call": + args = event.get("arguments", {}) + trace.print( + f" [dim]↳ Searching:[/dim] " + f"[dim italic]{_format_search_call(args)}[/dim italic]" + ) + elif etype == "search_result": + n = event.get("num_hits", 0) + label = "result" if n == 1 else "results" + trace.print(f" [dim]↳ Found {n} {label}[/dim]") + elif etype == "clarify_call": + q = event.get("question", "") or "" + trace.print( + f" [dim]↳ Clarifying:[/dim] [dim italic]{q}[/dim italic]" + ) + # final_answer and clarify_response are handled outside the loop. + + agent = ResearchAgent( + engine=engine, + search=HybridSearch(store, embedder), + model=planner_model, + on_event=on_event, + ) + + started = time.monotonic() + result = agent.run(query_text) + elapsed = time.monotonic() - started + + logger.debug( + "research: iterations=%d tool_calls=%d usage=%s", + result.iterations, + len(result.tool_calls), + result.usage, + ) + + if output_json: + click.echo( + json_mod.dumps( + { + "answer": result.answer, + "iterations": result.iterations, + "usage": result.usage, + "tool_calls": [ + { + "arguments": inv.arguments, + "num_results": inv.num_results, + "top_titles": inv.top_titles, + } + for inv in result.tool_calls + ], + }, + indent=2, + ) + ) + return + + # Visual break between live traces and the synthesis. + trace.print() + + # Render the synthesis as Markdown with inline citations coloured cyan. + answer_console.print(Markdown(_style_citations(result.answer))) + + # Footer: how long it took + how many distinct sources the model + # actually cited. Empty answers (rare; only if the model went silent) + # skip the footer entirely. + if result.answer: + cited = {int(n) for n in re.findall(r"\[(\d+)\]", result.answer)} + src_word = "source" if len(cited) == 1 else "sources" + trace.print() + trace.print( + f"[dim]Deep Research · {elapsed:.1f}s · " + f"{len(cited)} {src_word} cited[/dim]" + ) + + def _get_memory_backend(config): """Try to instantiate the memory backend. @@ -368,6 +549,24 @@ def _print_profile( is_flag=True, help="Print inference telemetry profile (latency, tokens, energy, IPW).", ) +@click.option( + "--research", + "research_mode", + is_flag=True, + help=( + "Route the query through the hybrid-search research agent over the " + "personal knowledge store (BM25 + dense embeddings, max 5 tool calls)." + ), +) +@click.option( + "--knowledge-db", + "knowledge_db", + default=None, + help=( + "Override the KnowledgeStore path used by --research " + "(default: ~/.openjarvis/knowledge.db)." + ), +) def ask( query: tuple[str, ...], model_name: str | None, @@ -380,6 +579,8 @@ def ask( agent_name: str | None, tool_names: str | None, enable_profile: bool, + research_mode: bool, + knowledge_db: str | None, ) -> None: """Ask Jarvis a question.""" console = Console(stderr=True) @@ -445,6 +646,20 @@ def ask( engine_name, engine = resolved + # ------------------------------------------------------------------ + # Research mode — hybrid search + agentic loop over the knowledge store + # ------------------------------------------------------------------ + if research_mode: + _run_research( + query_text=query_text, + engine=engine, + model_name=model_name, + knowledge_db=knowledge_db, + output_json=output_json, + console=console, + ) + return + # Apply security guardrails from openjarvis.security import setup_security diff --git a/src/openjarvis/connectors/_stubs.py b/src/openjarvis/connectors/_stubs.py index 6ea4e95a..e956bfe8 100644 --- a/src/openjarvis/connectors/_stubs.py +++ b/src/openjarvis/connectors/_stubs.py @@ -26,6 +26,11 @@ class Document: """Universal schema for data from any connector. All connectors normalize their output to this format before ingestion. + + v1 schema fields (``source_id``, ``participants_raw``, ``channel``) default + to empty so existing connectors compile without modification; new + connectors should populate them. The pipeline derives ``source_id`` from + ``doc_id`` by stripping the ``{source}:`` prefix when not set explicitly. """ doc_id: str @@ -40,6 +45,10 @@ class Document: url: Optional[str] = None attachments: List[Attachment] = field(default_factory=list) metadata: Dict[str, Any] = field(default_factory=dict) + # v1 schema additions: defaulted empty so legacy connectors keep working. + source_id: str = "" + participants_raw: List[str] = field(default_factory=list) + channel: Optional[str] = None @dataclass(slots=True) diff --git a/src/openjarvis/connectors/chunker.py b/src/openjarvis/connectors/chunker.py index f0515497..b81c6131 100644 --- a/src/openjarvis/connectors/chunker.py +++ b/src/openjarvis/connectors/chunker.py @@ -1,20 +1,22 @@ """Type-aware semantic chunker for Deep Research ingestion. -Splits text based on document type, never splitting mid-sentence. -Returns ``ChunkResult`` dataclass objects with section metadata and -inherited parent metadata. +Splits text by paragraph → sentence → token/character boundaries while +enforcing a hard size cap and adding a fixed-size overlap between +consecutive chunks. The cap is enforced on BOTH token count +(whitespace-split) AND character count, since marketing-style emails +frequently contain dense runs without whitespace (zero-width joiners, +HTML residue) that defeat token-based limits alone. Splitting strategy by doc_type ------------------------------- -- ``event``, ``contact`` : Always a single chunk; never split. +- ``event``, ``contact`` : Always a single chunk; never split, never capped. - ``email`` : Split on reply boundaries (``On … wrote:``), - then sentence-split within each part. -- ``message`` : Split on double-newline boundaries, accumulate - into chunks up to *max_tokens*. + then paragraphs, then sentences, then force-split. +- ``message`` : Split on double-newline boundaries, then sentences, + then force-split. - ``document``, ``note``, anything else : Split on ``## Heading`` section boundaries → - paragraph boundaries (``\\n\\n``) within sections → - sentence boundaries as a last resort. + paragraph boundaries → sentences → force-split. Token counting uses whitespace splitting: ``len(text.split())``. """ @@ -23,15 +25,21 @@ from __future__ import annotations import re from dataclasses import dataclass, field -from typing import Any, Dict, List, Optional +from typing import Any, Dict, List, Optional, Tuple # --------------------------------------------------------------------------- -# Public types +# Regexes # --------------------------------------------------------------------------- _SENTENCE_SPLIT_RE = re.compile(r'(?<=[.!?])\s+(?=[A-Z"])') _SECTION_RE = re.compile(r"(?m)^##\s+(.+)$") _REPLY_BOUNDARY_RE = re.compile(r"(?m)^On .+wrote:\s*$") +_SENTENCE_END_RE = re.compile(r"[.!?](?:\s|$)") + + +# --------------------------------------------------------------------------- +# Public types +# --------------------------------------------------------------------------- @dataclass(slots=True) @@ -54,65 +62,118 @@ def _count_tokens(text: str) -> int: def _split_sentences(text: str) -> List[str]: - """Split *text* into sentences using the canonical regex. - - The regex splits after sentence-ending punctuation (``.``, ``!``, ``?``) - followed by whitespace and a capital letter or a double-quote. - """ parts = _SENTENCE_SPLIT_RE.split(text) return [p.strip() for p in parts if p.strip()] -def _accumulate( +def _accumulate_capped( segments: List[str], *, max_tokens: int, + max_chars: int, sep: str = " ", ) -> List[str]: - """Greedily merge *segments* into chunks up to *max_tokens* tokens. + """Greedy-merge segments into chunks bounded by token AND char counts. - A segment that is already larger than *max_tokens* is placed in its own - chunk; it is never split further by this function. + Segments larger than the bounds pass through as their own chunk — + the caller force-splits those in a separate pass. """ chunks: List[str] = [] - current_parts: List[str] = [] - current_tokens = 0 + current: List[str] = [] + cur_tokens = 0 + cur_chars = 0 for seg in segments: seg_tokens = _count_tokens(seg) - if current_parts and current_tokens + seg_tokens > max_tokens: - chunks.append(sep.join(current_parts)) - current_parts = [seg] - current_tokens = seg_tokens + seg_chars = len(seg) + sep_chars = len(sep) if current else 0 + + would_overflow = current and ( + cur_tokens + seg_tokens > max_tokens + or cur_chars + sep_chars + seg_chars > max_chars + ) + if would_overflow: + chunks.append(sep.join(current)) + current = [seg] + cur_tokens = seg_tokens + cur_chars = seg_chars else: - current_parts.append(seg) - current_tokens += seg_tokens - - if current_parts: - chunks.append(sep.join(current_parts)) + current.append(seg) + cur_tokens += seg_tokens + cur_chars += sep_chars + seg_chars + if current: + chunks.append(sep.join(current)) return chunks -def _sentence_chunks(text: str, *, max_tokens: int) -> List[str]: - """Split *text* by sentences and accumulate into max_tokens chunks.""" - sentences = _split_sentences(text) - if not sentences: - stripped = text.strip() - return [stripped] if stripped else [] - return _accumulate(sentences, max_tokens=max_tokens, sep=" ") +def _best_cut_index(window: str) -> int: + """Pick a cut point inside ``window``: sentence end → space → hard cut.""" + matches = list(_SENTENCE_END_RE.finditer(window)) + if matches: + return matches[-1].end() + midpoint = len(window) // 2 + space = window.rfind(" ", midpoint) + if space > 0: + return space + 1 + return len(window) -def _paragraph_chunks(text: str, *, max_tokens: int) -> List[str]: - """Split *text* on paragraph breaks (``\\n\\n``), then by sentences if needed.""" - paragraphs = [p.strip() for p in text.split("\n\n") if p.strip()] - result: List[str] = [] - for para in paragraphs: - if _count_tokens(para) <= max_tokens: - result.append(para) - else: - result.extend(_sentence_chunks(para, max_tokens=max_tokens)) - return result +def _force_split(text: str, *, max_chars: int, max_tokens: int) -> List[str]: + """Split an oversized run into pieces that fit both caps. + + Walks the text greedily: takes a window up to ``max_chars``, shrinks + until the token count fits, then snaps the cut to the last sentence + boundary, falling back to a word boundary, then a hard cut. + """ + rest = text.strip() + out: List[str] = [] + while rest: + if len(rest) <= max_chars and _count_tokens(rest) <= max_tokens: + out.append(rest) + break + + end = min(max_chars, len(rest)) + while end > 1 and _count_tokens(rest[:end]) > max_tokens: + end = max(1, int(end * 0.9)) + + window = rest[:end] + cut = _best_cut_index(window) + piece = rest[:cut].strip() + if not piece: + # Window contained only whitespace before any cuttable boundary. + # Force advance to avoid an infinite loop. + cut = max(cut + 1, end) + piece = rest[:cut].strip() + if piece: + out.append(piece) + rest = rest[cut:].strip() + return out + + +def _apply_overlap(chunks: List[str], *, overlap_tokens: int) -> List[str]: + """Prepend the last ``overlap_tokens`` tokens of each chunk to the next. + + The tail is taken from the *original* preceding chunk (not the + already-augmented output) so overlap doesn't compound, and is also + capped by character count — a single whitespace-split "token" can be + a 200+ char URL, so a pure-token cap would let tails balloon on + content with long opaque strings. + """ + if overlap_tokens <= 0 or len(chunks) < 2: + return list(chunks) + overlap_chars_cap = overlap_tokens * 4 + out = [chunks[0]] + for i in range(1, len(chunks)): + prev_tokens = chunks[i - 1].split() + if not prev_tokens: + out.append(chunks[i]) + continue + tail = " ".join(prev_tokens[-overlap_tokens:]) + if len(tail) > overlap_chars_cap: + tail = tail[-overlap_chars_cap:] + out.append(f"{tail} {chunks[i]}".strip()) + return out # --------------------------------------------------------------------------- @@ -121,18 +182,55 @@ def _paragraph_chunks(text: str, *, max_tokens: int) -> List[str]: class SemanticChunker: - """Split text based on document type without breaking mid-sentence. + """Split text by document type with a hard size cap and overlap. Parameters ---------- max_tokens: - Soft upper limit on chunk size measured in whitespace-delimited tokens - (i.e. ``len(text.split())``). Single unsplittable segments may exceed - this limit. + Hard upper bound on chunk size in whitespace-delimited tokens. + No emitted chunk exceeds this. + max_chars: + Hard upper bound on chunk size in characters. Defaults to + ``max_tokens * 4`` (a typical English chars-per-token estimate). + No emitted chunk exceeds this — the chunker force-splits any + run that does, even when token count alone would say it fits. + overlap_tokens: + Token tail copied from each chunk into the head of the next so + downstream retrieval doesn't miss context that straddles a chunk + boundary. Defaults to ``min(100, max_tokens // 5)`` and is clamped + to ``[0, max_tokens - 1]``. Set to ``0`` to disable. """ - def __init__(self, max_tokens: int = 512) -> None: + def __init__( + self, + max_tokens: int = 512, + *, + max_chars: Optional[int] = None, + overlap_tokens: Optional[int] = None, + ) -> None: self.max_tokens = max_tokens + self.max_chars = max_chars if max_chars is not None else max_tokens * 4 + if overlap_tokens is None: + overlap_tokens = min(100, max(1, max_tokens // 5)) + self.overlap_tokens = max(0, min(overlap_tokens, max(0, max_tokens - 1))) + + # Content budget per chunk leaves room for the overlap prefix + # that gets added in the final pass, so tokens stay within + # max_tokens. Char budget intentionally does NOT subtract overlap + # — we accept up to ~overlap_tokens*4 chars of headroom over + # max_chars after overlap, which still sits under any reasonable + # hard ceiling and avoids spurious force-splits of well-behaved + # sentences in configurations where the char cap is tight. + self._content_tokens = max(1, self.max_tokens - self.overlap_tokens) + self._content_chars = self.max_chars + # Soft target used to decide when a paragraph is "long enough to + # sub-split", and to size sentence-accumulated sub-chunks. At + # roughly half the hard cap, typical paragraphs ship as a single + # chunk and only the unusually long ones get further split, + # which keeps semantic units intact while still pulling the + # chunk-length tail down. + self._target_tokens = max(1, self._content_tokens // 2) + self._target_chars = max(1, self._content_chars // 2) # ------------------------------------------------------------------ # Public API @@ -147,16 +245,10 @@ class SemanticChunker: ) -> List[ChunkResult]: """Split *text* into ``ChunkResult`` objects. - Parameters - ---------- - text: The raw text to split. - doc_type: Controls the splitting strategy (see module docstring). - metadata: Parent metadata dict; copied into every chunk's ``metadata``. - - Returns - ------- - A list of ``ChunkResult`` objects with sequential 0-based ``index`` - values. Returns an empty list if *text* is empty or whitespace-only. + Returns an empty list if *text* is empty or whitespace-only. + Events and contacts are always returned as a single chunk + regardless of size; all other types respect the size caps and + receive overlap between consecutive chunks. """ if not text or not text.strip(): return [] @@ -164,88 +256,115 @@ class SemanticChunker: parent_meta: Dict[str, Any] = dict(metadata or {}) if doc_type in ("event", "contact"): - raw_chunks = self._chunk_atomic(text) - elif doc_type == "email": + return [ChunkResult(content=text, index=0, metadata=parent_meta)] + + if doc_type == "email": raw_chunks = self._chunk_email(text) elif doc_type == "message": raw_chunks = self._chunk_message(text) else: - # "document", "note", or any unknown type raw_chunks = self._chunk_document(text) + # Apply overlap globally between consecutive chunks. + contents = [c for c, _ in raw_chunks] + metas = [m for _, m in raw_chunks] + overlapped = _apply_overlap(contents, overlap_tokens=self.overlap_tokens) + results: List[ChunkResult] = [] - for idx, (content, extra_meta) in enumerate(raw_chunks): + for idx, (content, extra_meta) in enumerate(zip(overlapped, metas)): merged: Dict[str, Any] = dict(parent_meta) merged.update(extra_meta) results.append(ChunkResult(content=content, index=idx, metadata=merged)) - return results # ------------------------------------------------------------------ # Strategy implementations # ------------------------------------------------------------------ - def _chunk_atomic(self, text: str) -> List[tuple[str, Dict[str, Any]]]: - """Return the entire text as a single chunk (event / contact).""" - return [(text, {})] + def _pack_text(self, text: str) -> List[str]: + """Emit one chunk per paragraph; sub-split paragraphs over the soft target. - def _chunk_email(self, text: str) -> List[tuple[str, Dict[str, Any]]]: - """Split on reply boundaries; sentence-split each part.""" - # Split the email into parts on "On ... wrote:" lines. - # re.split with a capturing group keeps the boundary in results, - # so we re-attach the header to the following segment. + Paragraphs are the natural chunk unit. A paragraph that fits the + soft target ships as one chunk. A paragraph that doesn't is + sentence-split and the sentences accumulated up to the soft + target; sentences that exceed the hard cap (rare — opaque content + with no whitespace) are force-split. + """ + paragraphs = [p.strip() for p in text.split("\n\n") if p.strip()] + if not paragraphs: + stripped = text.strip() + return [stripped] if stripped else [] + + out: List[str] = [] + for para in paragraphs: + if ( + len(para) <= self._target_chars + and _count_tokens(para) <= self._target_tokens + ): + out.append(para) + continue + + sents = _split_sentences(para) + if not sents: + sents = [para] + + accumulated = _accumulate_capped( + sents, + max_tokens=self._target_tokens, + max_chars=self._target_chars, + sep=" ", + ) + for c in accumulated: + if ( + len(c) <= self._content_chars + and _count_tokens(c) <= self._content_tokens + ): + out.append(c) + else: + out.extend( + _force_split( + c, + max_chars=self._content_chars, + max_tokens=self._content_tokens, + ) + ) + return [c for c in out if c] + + def _chunk_email(self, text: str) -> List[Tuple[str, Dict[str, Any]]]: + """Split on reply boundaries; pack each part.""" boundaries = _REPLY_BOUNDARY_RE.split(text) + headers = _REPLY_BOUNDARY_RE.findall(text) - # Each boundary match is a separator; reassemble so the "On … wrote:" - # line stays with the content that follows it (the quoted block). raw_parts: List[str] = [] if boundaries: - # The first element is the text before the first boundary (the - # main reply body). raw_parts.append(boundaries[0]) - # Subsequent elements alternate: matched boundary, then text after. - # Because we used split() (not findall), the boundaries themselves - # are not in the list — only the text segments between them. - # So boundaries[1:] are the segments after each matched header. - # We need to re-find the headers to reassemble. - headers = _REPLY_BOUNDARY_RE.findall(text) for header, body in zip(headers, boundaries[1:]): - # We found the header text via findall; reconstruct the part. part = (header.strip() + "\n" + body).strip() raw_parts.append(part) - chunks: List[tuple[str, Dict[str, Any]]] = [] + chunks: List[Tuple[str, Dict[str, Any]]] = [] for part in raw_parts: part = part.strip() if not part: continue - if _count_tokens(part) <= self.max_tokens: - chunks.append((part, {})) - else: - for sub in _sentence_chunks(part, max_tokens=self.max_tokens): - if sub: - chunks.append((sub, {})) + for c in self._pack_text(part): + if c: + chunks.append((c, {})) return chunks if chunks else [(text.strip(), {})] - def _chunk_message(self, text: str) -> List[tuple[str, Dict[str, Any]]]: - """Split on double-newline boundaries and accumulate up to max_tokens.""" - paragraphs = [p.strip() for p in text.split("\n\n") if p.strip()] - raw_chunks = _accumulate(paragraphs, max_tokens=self.max_tokens, sep="\n\n") - return [(c, {}) for c in raw_chunks if c] + def _chunk_message(self, text: str) -> List[Tuple[str, Dict[str, Any]]]: + """Pack a message via paragraphs → sentences → force-split.""" + return [(c, {}) for c in self._pack_text(text)] - def _chunk_document(self, text: str) -> List[tuple[str, Dict[str, Any]]]: - """Split on ## headings → paragraphs → sentences.""" - # Find all ## heading positions + def _chunk_document(self, text: str) -> List[Tuple[str, Dict[str, Any]]]: + """Split on ## headings → paragraphs → sentences → force-split.""" section_matches = list(_SECTION_RE.finditer(text)) if not section_matches: - # No headings — fall back to paragraph/sentence splitting - raw_chunks = _paragraph_chunks(text, max_tokens=self.max_tokens) - return [(c, {}) for c in raw_chunks if c] + return [(c, {}) for c in self._pack_text(text)] - # Build (title, body_text) pairs for each section - sections: List[tuple[str, str]] = [] + sections: List[Tuple[str, str]] = [] for i, m in enumerate(section_matches): title = m.group(1).strip() body_start = m.end() @@ -257,24 +376,19 @@ class SemanticChunker: body = text[body_start:body_end].strip() sections.append((title, body)) - # Check for preamble text before the first heading + result: List[Tuple[str, Dict[str, Any]]] = [] preamble = text[: section_matches[0].start()].strip() - result: List[tuple[str, Dict[str, Any]]] = [] - if preamble: - for c in _paragraph_chunks(preamble, max_tokens=self.max_tokens): + for c in self._pack_text(preamble): if c: result.append((c, {})) for title, body in sections: section_meta: Dict[str, Any] = {"section": title} if not body: - # Empty section — emit a placeholder chunk with just the title result.append((title, section_meta)) continue - - para_chunks = _paragraph_chunks(body, max_tokens=self.max_tokens) - for c in para_chunks: + for c in self._pack_text(body): if c: result.append((c, dict(section_meta))) diff --git a/src/openjarvis/connectors/embeddings.py b/src/openjarvis/connectors/embeddings.py new file mode 100644 index 00000000..45b35119 --- /dev/null +++ b/src/openjarvis/connectors/embeddings.py @@ -0,0 +1,156 @@ +"""Dense embedding clients for the IngestionPipeline. + +A thin HTTP wrapper around a local `Ollama `_ daemon +running an embedding model (default ``nomic-embed-text``, 768-dim). Embeddings +are serialised as float32 ``bytes`` for storage in the ``embedding`` BLOB +column of ``knowledge_chunks``. + +The client degrades gracefully when the daemon is unreachable: ``embed`` +returns ``None`` and ``is_available()`` reports ``False`` instead of raising, +so ingestion never fails because a sidecar service is down. +""" + +from __future__ import annotations + +import logging +from typing import List, Optional + +import numpy as np +import requests + +logger = logging.getLogger(__name__) + + +DEFAULT_OLLAMA_HOST = "http://localhost:11434" +DEFAULT_EMBED_MODEL = "nomic-embed-text" + + +class OllamaEmbedder: + """Embed text via a local Ollama daemon. + + Parameters + ---------- + model: + Ollama model tag (e.g. ``nomic-embed-text``, ``mxbai-embed-large``). + host: + Base URL for the Ollama HTTP API. Defaults to ``http://localhost:11434``. + timeout: + Per-request timeout in seconds. + """ + + def __init__( + self, + *, + model: str = DEFAULT_EMBED_MODEL, + host: str = DEFAULT_OLLAMA_HOST, + timeout: float = 30.0, + ) -> None: + self._model = model + self._host = host.rstrip("/") + self._timeout = timeout + self._dim: Optional[int] = None + + # ------------------------------------------------------------------ + # Identity / capability checks + # ------------------------------------------------------------------ + + @property + def model_version(self) -> str: + """Stable identifier persisted alongside each embedding row.""" + return f"ollama:{self._model}" + + @property + def dim(self) -> Optional[int]: + """Embedding dimensionality, learned after the first successful call.""" + return self._dim + + def is_available(self) -> bool: + """Return True iff the daemon answers and the model is installed.""" + try: + resp = requests.get(f"{self._host}/api/tags", timeout=2.0) + resp.raise_for_status() + except requests.RequestException: + return False + + try: + names = {m.get("name", "") for m in resp.json().get("models", [])} + except ValueError: + return False + + # Ollama tags include the ":latest" suffix; match either form. + return self._model in names or f"{self._model}:latest" in names + + # ------------------------------------------------------------------ + # Embedding + # ------------------------------------------------------------------ + + def embed(self, text: str) -> Optional[bytes]: + """Embed a single string. Returns float32 bytes or ``None`` on failure.""" + if not text or not text.strip(): + return None + try: + resp = requests.post( + f"{self._host}/api/embeddings", + json={"model": self._model, "prompt": text}, + timeout=self._timeout, + ) + resp.raise_for_status() + payload = resp.json() + except requests.RequestException as exc: + logger.warning("OllamaEmbedder.embed: request failed (%s)", exc) + return None + except ValueError as exc: + logger.warning("OllamaEmbedder.embed: bad JSON (%s)", exc) + return None + + vec = payload.get("embedding") + if not vec: + logger.warning( + "OllamaEmbedder.embed: empty embedding for %d chars", len(text) + ) + return None + + arr = np.asarray(vec, dtype=np.float32) + if self._dim is None: + self._dim = int(arr.shape[0]) + elif arr.shape[0] != self._dim: + logger.warning( + "OllamaEmbedder.embed: dim drift (expected %d, got %d)", + self._dim, arr.shape[0], + ) + return None + return arr.tobytes() + + def embed_batch(self, texts: List[str]) -> List[Optional[bytes]]: + """Embed a list of strings sequentially. + + Ollama's HTTP API serves one prompt per call; on the same host the + round-trip overhead is negligible relative to model inference. + """ + return [self.embed(t) for t in texts] + + +# --------------------------------------------------------------------------- +# Deserialisation helper (used by verification + future retrieval code) +# --------------------------------------------------------------------------- + + +def decode_embedding( + blob: Optional[bytes], *, dtype: type = np.float32 +) -> Optional[np.ndarray]: + """Reconstruct a 1-D vector from a BLOB written by ``OllamaEmbedder.embed``. + + Returns ``None`` when the input is missing or zero-length so callers can + treat absent embeddings uniformly. + """ + if not blob: + return None + return np.frombuffer(blob, dtype=dtype) + + +__all__ = [ + "OllamaEmbedder", + "decode_embedding", + "DEFAULT_EMBED_MODEL", + "DEFAULT_OLLAMA_HOST", +] diff --git a/src/openjarvis/connectors/gcalendar.py b/src/openjarvis/connectors/gcalendar.py index c2548812..1a0b7735 100644 --- a/src/openjarvis/connectors/gcalendar.py +++ b/src/openjarvis/connectors/gcalendar.py @@ -13,6 +13,7 @@ from typing import Any, Dict, Iterator, List, Optional import httpx from openjarvis.connectors._stubs import BaseConnector, Document, SyncStatus +from openjarvis.connectors.google_auth import call_with_refresh from openjarvis.connectors.oauth import ( GOOGLE_ALL_SCOPES, build_google_auth_url, @@ -302,13 +303,15 @@ class GCalendarConnector(BaseConnector): tokens = load_tokens(self._credentials_path) if not tokens: return - - token: str = tokens.get("access_token", tokens.get("token", "")) - if not token: + if not tokens.get("access_token") and not tokens.get("token"): return - # Fetch list of calendars - calendars_resp = _gcal_api_calendars_list(token) + # Fetch list of calendars. call_with_refresh wraps the token read so + # an expired access_token triggers a one-shot refresh + retry instead + # of bubbling up a 401. + calendars_resp = call_with_refresh( + _gcal_api_calendars_list, self._credentials_path + ) calendars: List[Dict[str, Any]] = calendars_resp.get("items", []) # Default to 24 hours ago so we don't dump the entire calendar history @@ -327,8 +330,9 @@ class GCalendarConnector(BaseConnector): while True: try: - events_resp = _gcal_api_events_list( - token, + events_resp = call_with_refresh( + _gcal_api_events_list, + self._credentials_path, calendar_id, page_token=page_token, time_min=time_min, diff --git a/src/openjarvis/connectors/gcontacts.py b/src/openjarvis/connectors/gcontacts.py index 58852ba9..08659f0d 100644 --- a/src/openjarvis/connectors/gcontacts.py +++ b/src/openjarvis/connectors/gcontacts.py @@ -13,6 +13,7 @@ from typing import Any, Dict, Iterator, List, Optional import httpx from openjarvis.connectors._stubs import BaseConnector, Document, SyncStatus +from openjarvis.connectors.google_auth import call_with_refresh from openjarvis.connectors.oauth import ( GOOGLE_ALL_SCOPES, build_google_auth_url, @@ -249,16 +250,16 @@ class GContactsConnector(BaseConnector): tokens = load_tokens(self._credentials_path) if not tokens: return - - token: str = tokens.get("access_token", tokens.get("token", "")) - if not token: + if not tokens.get("access_token") and not tokens.get("token"): return page_token: Optional[str] = cursor synced = 0 while True: - list_resp = _gcontacts_api_list(token, page_token=page_token) + list_resp = call_with_refresh( + _gcontacts_api_list, self._credentials_path, page_token=page_token + ) connections: List[Dict[str, Any]] = list_resp.get("connections", []) for person in connections: diff --git a/src/openjarvis/connectors/gdrive.py b/src/openjarvis/connectors/gdrive.py index 18cf0455..5b7f4cf9 100644 --- a/src/openjarvis/connectors/gdrive.py +++ b/src/openjarvis/connectors/gdrive.py @@ -13,6 +13,7 @@ from typing import Any, Dict, Iterator, List, Optional import httpx from openjarvis.connectors._stubs import BaseConnector, Document, SyncStatus +from openjarvis.connectors.google_auth import call_with_refresh from openjarvis.connectors.oauth import ( GOOGLE_ALL_SCOPES, build_google_auth_url, @@ -234,16 +235,16 @@ class GDriveConnector(BaseConnector): tokens = load_tokens(self._credentials_path) if not tokens: return - - token: str = tokens.get("access_token", tokens.get("token", "")) - if not token: + if not tokens.get("access_token") and not tokens.get("token"): return page_token: Optional[str] = cursor synced = 0 while True: - list_resp = _gdrive_api_list_files(token, page_token=page_token) + list_resp = call_with_refresh( + _gdrive_api_list_files, self._credentials_path, page_token=page_token + ) files: List[Dict[str, Any]] = list_resp.get("files", []) for file_meta in files: @@ -263,7 +264,12 @@ class GDriveConnector(BaseConnector): export_mime = _EXPORT_MIME_MAP.get(mime_type) if export_mime is not None: try: - content = _gdrive_api_export(token, file_id, export_mime) + content = call_with_refresh( + _gdrive_api_export, + self._credentials_path, + file_id, + export_mime, + ) except Exception: # noqa: BLE001 content = f"[File: {name}] ({mime_type})" else: diff --git a/src/openjarvis/connectors/gmail.py b/src/openjarvis/connectors/gmail.py index ee157a83..64e87a87 100644 --- a/src/openjarvis/connectors/gmail.py +++ b/src/openjarvis/connectors/gmail.py @@ -9,12 +9,21 @@ from __future__ import annotations import base64 import email.utils +import logging +import re from datetime import datetime -from typing import Any, Dict, Iterator, List, Optional +from html.parser import HTMLParser +from typing import Any, Dict, Iterator, List, Optional, Tuple import httpx from openjarvis.connectors._stubs import BaseConnector, Document, SyncStatus +from openjarvis.connectors.google_auth import ( + GoogleAuthError, +) +from openjarvis.connectors.google_auth import ( + call_with_refresh as _call_with_refresh, +) from openjarvis.connectors.oauth import ( GOOGLE_ALL_SCOPES, build_google_auth_url, @@ -27,6 +36,8 @@ from openjarvis.core.config import DEFAULT_CONFIG_DIR from openjarvis.core.registry import ConnectorRegistry from openjarvis.tools._stubs import ToolSpec +logger = logging.getLogger(__name__) + # --------------------------------------------------------------------------- # Constants # --------------------------------------------------------------------------- @@ -35,6 +46,13 @@ _GMAIL_API_BASE = "https://gmail.googleapis.com/gmail/v1/users/me" _GMAIL_SCOPE = "https://www.googleapis.com/auth/gmail.readonly" _DEFAULT_CREDENTIALS_PATH = str(DEFAULT_CONFIG_DIR / "connectors" / "gmail.json") +# Token refresh is delegated to the shared google_auth helper so Calendar, +# Contacts, Drive, and Tasks get the same one-shot 401 retry. ``GmailAuthError`` +# is preserved as a module-level alias for tests/callers that imported it +# under the historical name. +GmailAuthError = GoogleAuthError + + # --------------------------------------------------------------------------- # Module-level API functions (easy to patch in tests) # --------------------------------------------------------------------------- @@ -118,21 +136,86 @@ def _extract_header(headers: List[Dict[str, str]], name: str) -> str: return "" +class _HTMLTextExtractor(HTMLParser): + """Strip HTML tags and return readable text using stdlib only. + + Skips +

Visible content here.

+ + + """ + text = _html_to_text(html) + assert "Visible content here" in text + assert "color: red" not in text + assert "alert" not in text + assert "Ignore me" not in text + + +def test_html_to_text_decodes_entities() -> None: + """Named and numeric HTML entities are decoded to their characters.""" + from openjarvis.connectors.gmail import _html_to_text # noqa: PLC0415 + + html = "

Tom & Jerry — 50 cents

" + text = _html_to_text(html) + assert "Tom & Jerry" in text + assert "—" in text # — + assert "50" in text and "cents" in text + + +def test_html_to_text_inserts_paragraph_breaks() -> None: + """Block-level tags produce newlines so the chunker sees structure.""" + from openjarvis.connectors.gmail import _html_to_text # noqa: PLC0415 + + html = "
line one
line two
line three
" + text = _html_to_text(html) + # Each block produces at least one newline boundary. + assert text.count("\n") >= 2 + + +@patch("openjarvis.connectors.gmail._gmail_api_list_messages") +@patch("openjarvis.connectors.gmail._gmail_api_get_message") +def test_sync_strips_html_when_no_text_plain( + mock_get, + mock_list, + connector, + tmp_path: Path, +) -> None: + """A multipart message with only text/html yields stripped plain text.""" + creds_path = Path(connector._credentials_path) + creds_path.write_text(json.dumps({"token": "fake-access-token"}), encoding="utf-8") + + html_bytes = ( + b"

Hello world!

" + b"

Second paragraph.

" + ) + html_b64 = base64.urlsafe_b64encode(html_bytes).decode().rstrip("=") + + msg_html = { + "id": "msg-html-1", + "threadId": "thread-html-1", + "labelIds": ["INBOX"], + "payload": { + "mimeType": "multipart/alternative", + "headers": [ + {"name": "From", "value": "marketer@example.com"}, + {"name": "To", "value": "me@example.com"}, + {"name": "Subject", "value": "Marketing"}, + {"name": "Date", "value": "Mon, 01 Jan 2024 10:00:00 +0000"}, + ], + "parts": [ + # No text/plain alternative — only text/html. + { + "mimeType": "text/html", + "body": {"data": html_b64}, + }, + ], + }, + } + + mock_list.return_value = {"messages": [{"id": "msg-html-1"}]} + mock_get.return_value = msg_html + + docs = list(connector.sync()) + assert len(docs) == 1 + content = docs[0].content + assert "Hello" in content + assert "world" in content + assert "Second paragraph" in content + assert "<" not in content and ">" not in content + + +@patch("openjarvis.connectors.gmail._gmail_api_list_messages") +@patch("openjarvis.connectors.gmail._gmail_api_get_message") +def test_sync_prefers_text_plain_over_text_html( + mock_get, + mock_list, + connector, + tmp_path: Path, +) -> None: + """When both alternatives are present, text/plain wins over text/html.""" + creds_path = Path(connector._credentials_path) + creds_path.write_text(json.dumps({"token": "fake-access-token"}), encoding="utf-8") + + plain_b64 = base64.urlsafe_b64encode( + b"Plain text version preferred." + ).decode().rstrip("=") + html_b64 = base64.urlsafe_b64encode( + b"

HTML version

" + ).decode().rstrip("=") + + msg_alt = { + "id": "msg-alt-1", + "threadId": "thread-alt-1", + "labelIds": ["INBOX"], + "payload": { + "mimeType": "multipart/alternative", + "headers": [ + {"name": "From", "value": "alice@example.com"}, + {"name": "To", "value": "me@example.com"}, + {"name": "Subject", "value": "Both alternatives"}, + {"name": "Date", "value": "Mon, 01 Jan 2024 10:00:00 +0000"}, + ], + "parts": [ + {"mimeType": "text/plain", "body": {"data": plain_b64}}, + {"mimeType": "text/html", "body": {"data": html_b64}}, + ], + }, + } + + mock_list.return_value = {"messages": [{"id": "msg-alt-1"}]} + mock_get.return_value = msg_alt + + docs = list(connector.sync()) + assert len(docs) == 1 + assert docs[0].content == "Plain text version preferred." + + +# --------------------------------------------------------------------------- +# Auto-refresh on 401 +# --------------------------------------------------------------------------- + + +class _FakeResponse: + """Minimal stand-in for httpx.Response used by the refresh test.""" + + def __init__( + self, *, status_code: int, json_data: dict | None = None, text: str = "" + ): + self.status_code = status_code + self._json = json_data or {} + self.text = text + + def json(self): + return self._json + + def raise_for_status(self): + if self.status_code >= 400: + import httpx as _httpx + raise _httpx.HTTPStatusError( + f"HTTP {self.status_code}", request=None, response=self + ) + + +def _write_full_creds(tmp_path: Path) -> str: + """Write a credentials file with a refresh_token and client credentials.""" + creds_path = tmp_path / "gmail.json" + creds_path.write_text( + json.dumps( + { + "access_token": "old-access-token", + "refresh_token": "stored-refresh-token", + "client_id": "client-id-abc", + "client_secret": "client-secret-xyz", + } + ), + encoding="utf-8", + ) + return str(creds_path) + + +def test_401_triggers_refresh_and_retries_with_new_token(tmp_path: Path) -> None: + """A 401 on a Gmail API call refreshes the token, persists it, and retries.""" + from openjarvis.connectors import gmail as gmail_mod + + creds_path = _write_full_creds(tmp_path) + + # The first httpx.get returns 401, the second returns 200. + get_calls: list[dict] = [] + + def fake_get(url, *, headers, params, timeout): + get_calls.append({"url": url, "headers": dict(headers), "params": dict(params)}) + if len(get_calls) == 1: + return _FakeResponse(status_code=401, text="unauthorized") + return _FakeResponse(status_code=200, json_data={"id": "msg-1", "payload": {}}) + + post_calls: list[dict] = [] + + def fake_post(url, *, data, timeout): + post_calls.append({"url": url, "data": dict(data)}) + return _FakeResponse( + status_code=200, + json_data={"access_token": "fresh-access-token", "expires_in": 3599}, + ) + + with patch.object(gmail_mod.httpx, "get", side_effect=fake_get), \ + patch.object(gmail_mod.httpx, "post", side_effect=fake_post): + result = gmail_mod._call_with_refresh( + gmail_mod._gmail_api_get_message, creds_path, "msg-1" + ) + + # The retried request must carry the fresh token. + assert len(get_calls) == 2 + assert get_calls[0]["headers"]["Authorization"] == "Bearer old-access-token" + assert get_calls[1]["headers"]["Authorization"] == "Bearer fresh-access-token" + + # The refresh hit Google's token endpoint with the stored refresh credentials. + assert len(post_calls) == 1 + assert post_calls[0]["url"] == "https://oauth2.googleapis.com/token" + assert post_calls[0]["data"] == { + "client_id": "client-id-abc", + "client_secret": "client-secret-xyz", + "refresh_token": "stored-refresh-token", + "grant_type": "refresh_token", + } + + # The new access_token is persisted to the credentials file. + on_disk = json.loads(Path(creds_path).read_text(encoding="utf-8")) + assert on_disk["access_token"] == "fresh-access-token" + assert on_disk["token"] == "fresh-access-token" # legacy key kept in sync + assert on_disk["refresh_token"] == "stored-refresh-token" + assert on_disk["expires_in"] == 3599 + + # And the caller saw the body of the retried request. + assert result == {"id": "msg-1", "payload": {}} + + +def test_non_401_status_is_not_refreshed(tmp_path: Path) -> None: + """A 500 from Gmail must propagate — only 401 should trigger refresh.""" + import httpx as _httpx + + from openjarvis.connectors import gmail as gmail_mod + + creds_path = _write_full_creds(tmp_path) + + def fake_get(url, *, headers, params, timeout): + return _FakeResponse(status_code=503, text="service unavailable") + + fake_post = patch.object(gmail_mod.httpx, "post", side_effect=AssertionError( + "_call_with_refresh must not refresh on non-401 status" + )) + + with patch.object(gmail_mod.httpx, "get", side_effect=fake_get), fake_post: + with pytest.raises(_httpx.HTTPStatusError): + gmail_mod._call_with_refresh( + gmail_mod._gmail_api_get_message, creds_path, "msg-1" + ) + + +def test_refresh_raises_when_refresh_token_missing(tmp_path: Path) -> None: + """Refresh aborts with a clear error when no refresh_token is stored.""" + from openjarvis.connectors import google_auth + + creds_path = tmp_path / "gmail.json" + creds_path.write_text( + json.dumps({"access_token": "stale", "client_id": "c", "client_secret": "s"}), + encoding="utf-8", + ) + + with pytest.raises(google_auth.GoogleAuthError, match="refresh_token"): + google_auth.refresh_access_token(str(creds_path)) + + +def test_sync_recovers_when_list_returns_401(tmp_path: Path) -> None: + """An end-to-end sync survives a 401 on messages.list without re-auth. + + Sets up a connector against a creds file with full refresh credentials, + has the first list call return 401, the refresh return a new token, and + the retried list call return one message stub which is then fetched + successfully. Verifies the connector yields the expected Document and + that the credentials file is rewritten with the fresh token. + """ + from openjarvis.connectors import gmail as gmail_mod + + creds_path = _write_full_creds(tmp_path) + connector = gmail_mod.GmailConnector(credentials_path=creds_path) + + list_calls: list[str] = [] + + def fake_get(url, *, headers, params, timeout): + auth = headers.get("Authorization", "") + if "/messages/" in url: + # The single get_message call after the retried list call. + return _FakeResponse( + status_code=200, + json_data={ + "id": "msg-99", + "threadId": "thread-99", + "labelIds": ["INBOX"], + "payload": { + "mimeType": "text/plain", + "headers": [ + {"name": "From", "value": "alice@example.com"}, + {"name": "To", "value": "me@example.com"}, + {"name": "Subject", "value": "Hi"}, + { + "name": "Date", + "value": "Mon, 01 Jan 2024 10:00:00 +0000", + }, + ], + "body": {"data": "SGVsbG8="}, + }, + }, + ) + # /messages list call + list_calls.append(auth) + if len(list_calls) == 1: + return _FakeResponse(status_code=401, text="unauthorized") + return _FakeResponse( + status_code=200, + json_data={"messages": [{"id": "msg-99"}]}, + ) + + def fake_post(url, *, data, timeout): + return _FakeResponse( + status_code=200, + json_data={"access_token": "fresh-token-after-401", "expires_in": 3599}, + ) + + with patch.object(gmail_mod.httpx, "get", side_effect=fake_get), \ + patch.object(gmail_mod.httpx, "post", side_effect=fake_post): + docs: List[Document] = list(connector.sync()) + + assert len(docs) == 1 + assert docs[0].title == "Hi" + assert docs[0].author == "alice@example.com" + # First list call carried the stale token; second carried the fresh one. + assert list_calls == ["Bearer old-access-token", "Bearer fresh-token-after-401"] + # Creds file persisted the new token. + on_disk = json.loads(Path(creds_path).read_text(encoding="utf-8")) + assert on_disk["access_token"] == "fresh-token-after-401" diff --git a/tests/connectors/test_integration.py b/tests/connectors/test_integration.py index d7489a8e..0e3adc0e 100644 --- a/tests/connectors/test_integration.py +++ b/tests/connectors/test_integration.py @@ -55,7 +55,7 @@ def test_full_pipeline_obsidian_to_search(vault: Path, tmp_path: Path) -> None: assert cp is not None # SyncEngine checkpoint records total chunks ingested (not document count). # The two vault files produce 6 chunks via SemanticChunker (note strategy). - assert cp["items_synced"] == 6 # 6 chunks from 2 md files + assert cp["items_synced"] == 6 # 4. Search via knowledge_search tool tool = KnowledgeSearchTool(store=store) diff --git a/tests/connectors/test_pipeline.py b/tests/connectors/test_pipeline.py index fa70efce..0ed74969 100644 --- a/tests/connectors/test_pipeline.py +++ b/tests/connectors/test_pipeline.py @@ -292,3 +292,174 @@ def test_ingest_chunk_count_return_value( assert n2 == 1 # only doc_c is new assert store.count() == 3 + + +# --------------------------------------------------------------------------- +# v1 schema tests — pipeline-side derivation +# --------------------------------------------------------------------------- + + +def test_thread_id_namespaced_at_pipeline( + pipeline: IngestionPipeline, store: KnowledgeStore +) -> None: + """Pipeline prefixes raw thread_id with '{source}:' so connectors can't forget.""" + doc = _make_doc( + doc_id="gmail:msg-tn1", + source="gmail", + thread_id="raw-thread-id", + content="Thread namespacing should happen here.", + ) + pipeline.ingest([doc]) + + rows = store._conn.execute( + "SELECT thread_id FROM knowledge_chunks" + ).fetchall() + assert len(rows) == 1 + assert rows[0][0] == "gmail:raw-thread-id" + + +def test_thread_id_namespacing_is_idempotent( + pipeline: IngestionPipeline, store: KnowledgeStore +) -> None: + """A connector that already namespaced thread_id is not double-prefixed.""" + doc = _make_doc( + doc_id="gmail:msg-tn2", + source="gmail", + thread_id="gmail:already-prefixed", + content="Idempotent namespacing keeps a single prefix.", + ) + pipeline.ingest([doc]) + + rows = store._conn.execute( + "SELECT thread_id FROM knowledge_chunks" + ).fetchall() + assert rows[0][0] == "gmail:already-prefixed" + + +def test_source_id_derived_from_doc_id_prefix( + pipeline: IngestionPipeline, store: KnowledgeStore +) -> None: + """When source_id isn't set on the Document, the pipeline strips '{source}:' + from doc_id so legacy connectors keep working.""" + doc = _make_doc( + doc_id="gmail:msg42", + source="gmail", + content="Legacy connector with composite doc_id.", + ) + pipeline.ingest([doc]) + + rows = store._conn.execute( + "SELECT source_id FROM knowledge_chunks" + ).fetchall() + assert rows[0][0] == "msg42" + + +def test_source_id_uses_explicit_field_when_set( + pipeline: IngestionPipeline, store: KnowledgeStore +) -> None: + """An explicit Document.source_id wins over any prefix-stripping heuristic.""" + doc = _make_doc( + doc_id="some-other-shape", + source="gmail", + content="Explicit source_id should be used verbatim.", + ) + doc.source_id = "explicit-src-id" + pipeline.ingest([doc]) + + rows = store._conn.execute( + "SELECT source_id FROM knowledge_chunks" + ).fetchall() + assert rows[0][0] == "explicit-src-id" + + +def test_content_hash_computed_per_chunk( + pipeline: IngestionPipeline, store: KnowledgeStore +) -> None: + """content_hash equals sha256(chunk.content) for a single-chunk document.""" + import hashlib as _hashlib + + content = "Hello world this is a test document." + doc = _make_doc(doc_id="doc:hash:1", content=content) + pipeline.ingest([doc]) + + rows = store._conn.execute( + "SELECT content, content_hash FROM knowledge_chunks" + ).fetchall() + assert len(rows) == 1 + assert rows[0][1] == _hashlib.sha256(rows[0][0].encode("utf-8")).hexdigest() + + +def test_last_synced_set_at_ingest( + pipeline: IngestionPipeline, store: KnowledgeStore +) -> None: + """last_synced is populated with the ingest time, not 0 default.""" + import time as _time + + before = _time.time() + pipeline.ingest([_make_doc(doc_id="doc:ls:1", content="Last synced check.")]) + after = _time.time() + + rows = store._conn.execute( + "SELECT last_synced FROM knowledge_chunks" + ).fetchall() + assert len(rows) == 1 + assert before <= rows[0][0] <= after + + +# --------------------------------------------------------------------------- +# Embedding wire-up +# --------------------------------------------------------------------------- + + +class _StubEmbedder: + """Deterministic in-test embedder that mimics the OllamaEmbedder surface.""" + + model_version = "stub:test-embedder" + + def __init__(self) -> None: + self.calls = 0 + + def embed(self, text: str): # type: ignore[no-untyped-def] + import numpy as _np + self.calls += 1 + # Map content to a stable 4-d float32 vector for assertion convenience. + h = abs(hash(text)) % 10_000 + return _np.asarray([h, h + 1, h + 2, h + 3], dtype=_np.float32).tobytes() + + +def test_pipeline_populates_embedding_when_embedder_provided( + store: KnowledgeStore, +) -> None: + """Pipeline writes float32 embedding bytes + model_version when embedder is set.""" + import numpy as _np + + embedder = _StubEmbedder() + pipeline = IngestionPipeline(store, embedder=embedder) + pipeline.ingest([_make_doc(doc_id="doc:emb:1", content="Short embeddable text.")]) + + rows = store._conn.execute( + "SELECT embedding, embedding_model_version FROM knowledge_chunks" + ).fetchall() + assert len(rows) == 1 + blob, version = rows[0] + assert blob is not None + arr = _np.frombuffer(blob, dtype=_np.float32) + assert arr.shape == (4,) + assert version == "stub:test-embedder" + assert embedder.calls == 1 + + +def test_pipeline_skips_embedding_when_no_embedder( + pipeline: IngestionPipeline, store: KnowledgeStore, +) -> None: + """Default pipeline leaves embedding NULL and embedding_model_version empty.""" + pipeline.ingest( + [_make_doc(doc_id="doc:emb:none", content="No embedder configured.")] + ) + + rows = store._conn.execute( + "SELECT embedding, embedding_model_version FROM knowledge_chunks" + ).fetchall() + assert len(rows) == 1 + assert rows[0][0] is None + assert rows[0][1] == "" diff --git a/tests/connectors/test_store.py b/tests/connectors/test_store.py index c7f2f354..4d927863 100644 --- a/tests/connectors/test_store.py +++ b/tests/connectors/test_store.py @@ -335,6 +335,112 @@ def test_memory_retrieve_event_emitted(tmp_path: Path) -> None: assert EventType.MEMORY_RETRIEVE in types +# --------------------------------------------------------------------------- +# v1 schema tests +# --------------------------------------------------------------------------- + + +def test_v1_columns_round_trip_via_metadata(ks: KnowledgeStore) -> None: + """source_id, channel, content_hash, embedding_model_version, participants_raw + survive the store/retrieve round-trip via the metadata payload.""" + _store( + ks, + content="Email about partnership review", + source="gmail", + source_id="msg-abc123", + channel="INBOX", + content_hash="deadbeefcafe", + embedding_model_version="text-embedding-3-small/v1", + participants_raw=["Alice Bose "], + ) + results = ks.retrieve("partnership review", top_k=1) + assert len(results) == 1 + meta = results[0].metadata + assert meta.get("source_id") == "msg-abc123" + assert meta.get("channel") == "INBOX" + assert meta.get("content_hash") == "deadbeefcafe" + assert meta.get("embedding_model_version") == "text-embedding-3-small/v1" + assert meta.get("participants_raw") == ["Alice Bose "] + + +def test_deleted_at_filters_retrieve(ks: KnowledgeStore) -> None: + """Rows with non-NULL deleted_at are excluded from retrieve().""" + import time as _time + + _store(ks, content="Live document about quarterly research") + _store(ks, content="Tombstoned document about quarterly research") + + ks._conn.execute( + "UPDATE knowledge_chunks SET deleted_at = ? " + "WHERE content LIKE 'Tombstoned%'", + (_time.time(),), + ) + ks._conn.commit() + + results = ks.retrieve("quarterly research", top_k=10) + assert len(results) == 1 + assert results[0].content.startswith("Live") + + +def test_unique_natural_key_constraint(ks: KnowledgeStore) -> None: + """Duplicate (source, source_id, chunk_index) is silently skipped. + + The store uses ``INSERT OR IGNORE`` so re-running a sync or replaying a + dogfood script over an already-populated store no longer crashes; the + original row stays put and the duplicate is dropped. + """ + first_id = _store( + ks, + content="First copy", + source="gmail", + source_id="msg-unique-1", + chunk_index=0, + ) + second_id = _store( + ks, + content="Duplicate copy", + source="gmail", + source_id="msg-unique-1", + chunk_index=0, + ) + # Same identity — the natural key collision returned the existing row's id. + assert second_id == first_id + # Content of the original row is preserved. + row = ks._conn.execute( + "SELECT content FROM knowledge_chunks WHERE id = ?", (first_id,) + ).fetchone() + assert row["content"] == "First copy" + # Only one row exists for this natural key. + count = ks._conn.execute( + "SELECT COUNT(*) FROM knowledge_chunks " + "WHERE source = 'gmail' AND source_id = 'msg-unique-1' AND chunk_index = 0" + ).fetchone()[0] + assert count == 1 + + +def test_unique_constraint_skipped_when_source_id_empty(ks: KnowledgeStore) -> None: + """Legacy rows with empty source_id can co-exist (partial index excludes them).""" + _store(ks, content="Legacy doc one", source="legacy", source_id="") + # Should not raise — partial index is WHERE source_id != '' + _store(ks, content="Legacy doc two", source="legacy", source_id="") + assert ks.count() == 2 + + +def test_embedding_blob_round_trip(ks: KnowledgeStore) -> None: + """A bytes embedding payload survives storage and reads back unchanged.""" + payload = b"\x00\x01\x02\x03\x04\xff\xfe\xfd" + _store( + ks, + content="Vector-bearing document for embedding round trip", + source="test", + embedding=payload, + ) + row = ks._conn.execute( + "SELECT embedding FROM knowledge_chunks LIMIT 1" + ).fetchone() + assert bytes(row[0]) == payload + + def test_context_manager_closes_connection(tmp_path: Path) -> None: """Used as a context manager, KnowledgeStore closes its connection on exit.""" import sqlite3 diff --git a/uv.lock b/uv.lock index 0f66752f..324023e5 100644 --- a/uv.lock +++ b/uv.lock @@ -4982,6 +4982,7 @@ dependencies = [ { name = "httpx" }, { name = "openai" }, { name = "posthog" }, + { name = "pynvml" }, { name = "python-telegram-bot" }, { name = "rich" }, { name = "tomli", marker = "python_full_version < '3.11'" }, @@ -5236,6 +5237,7 @@ requires-dist = [ { name = "pygemma", marker = "extra == 'inference-gemma'", specifier = ">=0.1.3" }, { name = "pymessenger", marker = "extra == 'channel-messenger'", specifier = ">=0.0.7" }, { name = "pynostr", marker = "extra == 'channel-nostr'", specifier = ">=0.6" }, + { name = "pynvml", specifier = ">=13.0.1" }, { name = "pynvml", marker = "extra == 'energy-all'", specifier = ">=12.0" }, { name = "pynvml", marker = "extra == 'gpu-metrics'", specifier = ">=12.0" }, { name = "pytest", marker = "extra == 'dev'", specifier = ">=8" },