feat: Deep Research — personal deep research over Gmail with hybrid retrieval and agentic synthesis (#354)

This commit is contained in:
Robby Manihani
2026-05-18 17:46:12 -07:00
committed by GitHub
parent b863cbb07b
commit 9af0dfe336
50 changed files with 5253 additions and 443 deletions
+2
View File
@@ -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
+3
View File
@@ -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/
+23 -8
View File
@@ -146,14 +146,29 @@ export function ChatArea() {
</div>
) : (
<div className="max-w-[var(--chat-max-width)] mx-auto px-4 py-6">
{messages.map((msg) => (
<MessageBubble key={msg.id} message={msg} />
))}
{streamState.isStreaming && streamState.content === '' && (
<div className="flex justify-start mb-4">
<StreamingDots phase={streamState.phase} />
</div>
)}
{messages.map((msg, i) => {
const isLastAssistant =
i === messages.length - 1 && msg.role === 'assistant';
return (
<MessageBubble
key={msg.id}
message={msg}
isLive={isLastAssistant && streamState.isStreaming}
/>
);
})}
{(() => {
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 (
<div className="flex justify-start mb-4">
<StreamingDots phase={streamState.phase} />
</div>
);
})()}
</div>
)}
</div>
+235 -8
View File
@@ -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<number, ResearchSource>();
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 (
<div className="px-4 pb-4 pt-2" style={{ maxWidth: 'var(--chat-max-width)', margin: '0 auto', width: '100%' }}>
<div className="mb-2 flex flex-col gap-1">
<div className="flex items-center gap-2">
<button
type="button"
onClick={() => setDeepResearch(!deepResearch)}
disabled={streamState.isStreaming}
aria-pressed={deepResearch}
className="inline-flex items-center gap-1.5 px-2.5 py-1 rounded-full text-xs transition-colors cursor-pointer disabled:cursor-default disabled:opacity-50"
style={{
background: deepResearch ? 'var(--color-accent-subtle)' : 'transparent',
border: `1px solid ${deepResearch ? 'var(--color-accent)' : 'var(--color-border)'}`,
color: deepResearch ? 'var(--color-accent)' : 'var(--color-text-tertiary)',
}}
title={deepResearch ? 'Deep Research: on' : 'Deep Research: off'}
>
<Search size={12} />
Deep Research
</button>
</div>
{deepResearch && corpusSync.syncing && corpusSync.itemsSynced > 0 && (
<div
className="text-[11px] leading-snug"
style={{ color: 'var(--color-text-tertiary)' }}
>
Searching over{' '}
<span key={corpusSync.itemsSynced} className="sync-bump" style={{ color: 'var(--color-text-secondary)' }}>
{corpusSync.itemsSynced.toLocaleString()}
</span>{' '}
items sync in progress, results will improve as more data is indexed.
</div>
)}
</div>
<div
className="flex items-center gap-2 rounded-2xl px-4 py-3 transition-shadow"
style={{
+35 -3
View File
@@ -8,6 +8,8 @@ import 'katex/dist/katex.min.css';
import { Copy, Check } from 'lucide-react';
import { AudioPlayer } from './AudioPlayer';
import { ToolCallCard } from './ToolCallCard';
import { ResearchTimeline } from './ResearchTimeline';
import { rehypeCitations } from '../../lib/rehype-citations';
import { XRayFooter } from './XRayFooter';
import type { ChatMessage } from '../../types';
@@ -19,6 +21,7 @@ function stripThinkTags(text: string): string {
interface Props {
message: ChatMessage;
isLive?: boolean;
}
function getTextContent(node: any): string {
@@ -97,7 +100,7 @@ function CopyMessageButton({ content }: { content: string }) {
);
}
export function MessageBubble({ message }: Props) {
export function MessageBubble({ message, isLive = false }: Props) {
const isUser = message.role === 'user';
if (isUser) {
@@ -121,8 +124,33 @@ export function MessageBubble({ message }: Props) {
const cleanContent = useMemo(() => 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, NonNullable<ChatMessage['researchSources']>[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 (
<div className="group mb-6">
{/* Deep Research timeline (steps + status) */}
{(message.isResearch || (message.researchTraces && message.researchTraces.length > 0)) && (
<ResearchTimeline
traces={message.researchTraces ?? []}
isLive={isLive}
hasContent={cleanContent.length > 0}
/>
)}
{/* Tool calls */}
{message.toolCalls && message.toolCalls.length > 0 && (
<div className="mb-3 flex flex-col gap-2">
@@ -140,7 +168,7 @@ export function MessageBubble({ message }: Props) {
<div className="prose max-w-none">
<ReactMarkdown
remarkPlugins={[remarkGfm, remarkMath]}
rehypePlugins={[[rehypeHighlight, { detect: true }], rehypeKatex]}
rehypePlugins={rehypePlugins}
components={{
pre: CodeBlockPre,
}}
@@ -154,7 +182,11 @@ export function MessageBubble({ message }: Props) {
<div className="flex items-center gap-2 mt-1.5">
<CopyMessageButton content={cleanContent} />
</div>
<XRayFooter usage={message.usage} telemetry={message.telemetry} />
<XRayFooter
usage={message.usage}
telemetry={message.telemetry}
isResearch={message.isResearch}
/>
</div>
);
}
@@ -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 (
<div
className="text-xs leading-relaxed animate-pulse"
style={{ color: 'var(--color-text-tertiary)' }}
>
{text}
</div>
);
}
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 (
<div className="relative pl-4">
{/* Step indicator dot, sitting on the vertical rail */}
<div
className="absolute left-0 top-[7px] w-1.5 h-1.5 rounded-full"
style={{
background: active
? 'var(--color-accent)'
: 'var(--color-text-tertiary)',
opacity: active ? 1 : 0.6,
transform: 'translateX(-3px)',
boxShadow: active ? '0 0 0 3px var(--color-accent-subtle)' : 'none',
transition: 'background 200ms, box-shadow 200ms',
}}
/>
<div
className="text-[10px] uppercase tracking-[0.08em] mb-0.5"
style={{
color: active ? 'var(--color-accent)' : 'var(--color-text-tertiary)',
}}
>
Search {index}
</div>
<div
className="text-sm"
style={{ color: 'var(--color-text)', fontWeight: 450 }}
>
{trace.query}
</div>
{meta.length > 0 && (
<div
className="text-[11px] mt-0.5"
style={{ color: 'var(--color-text-tertiary)' }}
>
{meta.join(' · ')}
</div>
)}
{active ? (
<div
className="mt-2 h-px overflow-hidden relative"
style={{ background: 'var(--color-border)' }}
>
<div
className="research-shimmer absolute inset-y-0 w-1/4"
style={{ background: 'var(--color-accent)' }}
/>
</div>
) : trace.numHits != null ? (
<div
className="text-[11px] mt-1"
style={{ color: 'var(--color-text-tertiary)' }}
>
{trace.numHits} {trace.numHits === 1 ? 'result' : 'results'}
</div>
) : null}
{trace.topTitles && trace.topTitles.length > 0 && (
<div
className="text-[11px] mt-0.5 truncate"
style={{ color: 'var(--color-text-tertiary)', opacity: 0.75 }}
>
{trace.topTitles.slice(0, 2).join(' · ')}
</div>
)}
</div>
);
}
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 (
<div className="mb-4">
<div className="relative pl-4">
<div
className="absolute left-0 top-[7px] w-1.5 h-1.5 rounded-full"
style={{
background: 'var(--color-accent)',
transform: 'translateX(-3px)',
boxShadow: '0 0 0 3px var(--color-accent-subtle)',
}}
/>
<StatusLine text="Analyzing query" />
</div>
</div>
);
}
const summary = summarizeTraces(traces);
const Chevron = expanded ? ChevronUp : ChevronDown;
return (
<div className="mb-4">
<button
type="button"
onClick={() => setExpanded((e) => !e)}
className="flex items-center gap-2 text-[11px] mb-2 cursor-pointer transition-colors"
style={{
color: 'var(--color-text-tertiary)',
background: 'transparent',
border: 'none',
padding: 0,
}}
title={expanded ? 'Collapse search trace' : 'Expand search trace'}
>
<span>{summary}</span>
<Chevron size={12} />
</button>
{expanded && (
<div className="relative">
<div
aria-hidden
className="absolute top-1 bottom-1 left-0 w-px"
style={{ background: 'var(--color-border)' }}
/>
<div className="flex flex-col gap-3">
{traces.map((t, i) => (
<TimelineStep
key={t.id}
index={i + 1}
trace={t}
isLive={isLive}
/>
))}
{showSynthesizing && (
<div className="relative pl-4">
<div
className="absolute left-0 top-[7px] w-1.5 h-1.5 rounded-full"
style={{
background: 'var(--color-accent)',
transform: 'translateX(-3px)',
boxShadow: '0 0 0 3px var(--color-accent-subtle)',
}}
/>
<StatusLine text="Synthesizing findings" />
</div>
)}
</div>
</div>
)}
</div>
);
}
+5 -2
View File
@@ -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<EnergyData | null>(null);
const [telemetry, setTelemetry] = useState<TelemetryStats | null>(null);
@@ -129,13 +130,15 @@ export function SystemPanel() {
<MiniStat
icon={Zap}
label="Power"
value={(energy?.avg_power_w ?? 0).toFixed(1)}
value={(liveEnergy?.power_w ?? energy?.avg_power_w ?? 0).toFixed(1)}
unit="W"
/>
<MiniStat
icon={Activity}
label="Energy"
value={((energy?.total_energy_j ?? 0) / 1000).toFixed(1)}
value={(
((liveEnergy?.energy_j ?? energy?.total_energy_j ?? 0) / 1000)
).toFixed(1)}
unit="kJ"
/>
</div>
+14 -5
View File
@@ -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})` : ''}` });
}
+8 -2
View File
@@ -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() {
<Cpu size={14} />
)}
<div className="flex-1 min-w-0">
<span className="truncate block text-left" style={{ color: 'var(--color-text)' }}>
{selectedModel || serverInfo?.model || 'Select model'}
<span
className="truncate block text-left"
style={{ color: deepResearch ? 'var(--color-accent)' : 'var(--color-text)' }}
>
{deepResearch
? 'Deep Research'
: selectedModel || serverInfo?.model || 'Select model'}
</span>
{modelLoading && (
<span className="text-[10px] block text-left" style={{ color: 'var(--color-accent)' }}>
+46
View File
@@ -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;
}
+20 -6
View File
@@ -192,12 +192,26 @@ export async function checkHealth(): Promise<boolean> {
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<boolean> => {
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<unknown> {
+203
View File
@@ -0,0 +1,203 @@
import type { ResearchSource } from '../types';
interface HastNode {
type: string;
tagName?: string;
value?: string;
properties?: Record<string, unknown>;
children?: HastNode[];
}
// Tags whose descendants should NOT be searched for citations.
// `<a>` 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<number, ResearchSource>;
}) {
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)` → <a href="url">1</a>. 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 <a>/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);
};
}
+50 -1
View File
@@ -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<ResearchEvent> {
// /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();
}
}
+44
View File
@@ -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<StreamState>) => 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<Settings>) => void;
@@ -387,6 +403,8 @@ export const useAppStore = create<AppState>((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<AppState>((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<AppState>((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<AppState>((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 }),
+225 -113
View File
@@ -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 <Ico size={size} />;
};
// 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 (
<div style={{ marginTop: 12 }}>
<button
type="button"
onClick={() => setOpen((o) => !o)}
style={{
background: 'transparent',
border: 'none',
padding: 0,
fontSize: 11,
color: 'var(--color-text-tertiary)',
cursor: 'pointer',
textDecoration: 'underline',
}}
>
{open ? 'Hide advanced' : 'Advanced: Connect with Google OAuth'}
</button>
{open && (
<div
style={{
marginTop: 8,
padding: 10,
background: 'var(--color-bg)',
border: '1px solid var(--color-border)',
borderRadius: 6,
}}
>
<div style={{ fontSize: 11, color: 'var(--color-text-tertiary)', marginBottom: 8 }}>
For developers with an existing Google Cloud project. Enable the
Gmail API and create a Desktop OAuth client at{' '}
<a
href="https://console.cloud.google.com/apis/credentials"
target="_blank"
rel="noopener noreferrer"
style={{ color: 'var(--color-accent)', textDecoration: 'underline' }}
>
Google Cloud Credentials
</a>{' '}
then paste the Client ID and Client Secret below.
</div>
<InlineConnectForm
fields={[
{ name: 'email', placeholder: 'Client ID', type: 'text' },
{ name: 'password', placeholder: 'Client Secret', type: 'password' },
]}
loading={loading}
onSubmit={onConnect}
/>
</div>
)}
</div>
);
}
// ---------------------------------------------------------------------------
// 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 (
<div>
<div style={{ fontSize: 11, color: 'var(--color-warning)', marginBottom: 4 }}>
Indexed{' '}
<span key={totalIndexed} className="sync-bump">
{totalIndexed.toLocaleString()} {unitLabel}
</span>{' '}
<span style={{ color: 'var(--color-text-tertiary)' }}>
({rangeLabel})
</span>{' '}
<span style={{ color: 'var(--color-text-tertiary)' }}>
· Still indexing
</span>
</div>
<div style={{ fontSize: 10.5, color: 'var(--color-text-tertiary)' }}>
Deep Research available now · results improve as more {unitLabel} are indexed
</div>
</div>
);
}
// 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 (
<div>
<div style={{ display: 'flex', alignItems: 'center', gap: 8 }}>
<span style={{ fontSize: 12, color: 'var(--color-success)' }}>
{chunks.toLocaleString()} {unitLabel}
Indexed {totalIndexed.toLocaleString()} {unitLabel}
{rangeLabel && (
<span style={{ color: 'var(--color-text-tertiary)' }}>
{' '}({rangeLabel})
</span>
)}
{lastSyncLabel && (
<span style={{ color: 'var(--color-text-tertiary)' }}>
{' · '}Last synced {lastSyncLabel}
</span>
)}
</span>
<button
onClick={handleSync}
@@ -397,70 +558,14 @@ function SyncStatusDisplay({
);
}
// Actively syncing
if (sync?.state === 'syncing' || syncing) {
const pct = sync?.items_total && sync.items_total > 0
? Math.round((sync.items_synced / sync.items_total) * 100)
: null;
const label = sync?.items_total && sync.items_total > 0
? `${sync.items_synced.toLocaleString()} / ${sync.items_total.toLocaleString()}`
: sync?.items_synced && sync.items_synced > 0
? `${sync.items_synced.toLocaleString()} items so far`
: 'Starting...';
return (
<div>
<div style={{ fontSize: 11, color: 'var(--color-warning)', marginBottom: 4 }}>
Syncing {label}
</div>
<div style={{
height: 4, borderRadius: 2,
background: 'var(--color-bg-tertiary)',
overflow: 'hidden',
}}>
<div style={{
height: '100%', borderRadius: 2,
background: 'var(--color-warning)',
width: pct != null ? `${pct}%` : '30%',
transition: 'width 0.5s ease',
animationName: pct == null ? 'pulse' : undefined,
animationDuration: pct == null ? '1.5s' : undefined,
animationIterationCount: pct == null ? 'infinite' : undefined,
}} />
</div>
</div>
);
}
// Idle with items synced but no chunks yet (indexing)
if (sync?.state === 'idle' && sync.items_synced > 0) {
return (
<div>
<div style={{ fontSize: 11, color: 'var(--color-warning)', marginBottom: 4 }}>
Indexing {sync.items_synced.toLocaleString()} items...
</div>
<div style={{
height: 4, borderRadius: 2,
background: 'var(--color-bg-tertiary)',
overflow: 'hidden',
}}>
<div style={{
height: '100%', borderRadius: 2, background: 'var(--color-warning)',
width: '60%',
animationName: 'pulse', animationDuration: '1.5s', animationIterationCount: 'infinite',
}} />
</div>
</div>
);
}
// Connected but no chunks yet
// Connected but nothing ever ingested. Mirror the original copy.
const hasSynced = sync?.last_sync != null;
return (
<div>
<div style={{ display: 'flex', alignItems: 'center', gap: 8 }}>
<span style={{ fontSize: 12, color: 'var(--color-text-tertiary)' }}>
{hasSynced
? 'Synced — 0 items found'
? `Synced — 0 ${unitLabel} found`
: 'Connected — not synced yet'}
</span>
<button
@@ -546,6 +651,21 @@ function DataSourcesSection() {
const [connectingId, setConnectingId] = useState<string | null>(null);
const [connectStage, setConnectStage] = useState<string>('');
const [connectError, setConnectError] = useState<string>('');
const [disconnectingId, setDisconnectingId] = useState<string | null>(null);
const handleDisconnect = async (id: string) => {
if (disconnectingId) return;
setDisconnectingId(id);
try {
await disconnectSource(id);
loadConnectors();
} catch {
// Surface failures silently — the connector list will refresh on the
// next poll and reflect the true state regardless.
} finally {
setDisconnectingId(null);
}
};
const handleConnect = async (id: string, req: ConnectRequest) => {
setLoading(true);
@@ -598,8 +718,32 @@ function DataSourcesSection() {
}
};
const connected = connectors.filter((c) => c.connected);
const notConnectedBase = connectors.filter((c) => !c.connected);
// Merge the OAuth Gmail (`gmail`) and IMAP Gmail (`gmail_imap`) backend
// connectors into a single user-facing Gmail card. IMAP is the default
// flow (no Google Cloud setup needed); OAuth lives behind an "Advanced"
// disclosure when the card is expanded. If both happen to be connected,
// keep whichever has more indexed chunks so the active source still
// surfaces its sync state.
const unifiedConnectors = (() => {
const gmail = connectors.find((c) => c.connector_id === 'gmail');
const gmailImap = connectors.find((c) => c.connector_id === 'gmail_imap');
if (!gmail || !gmailImap) return connectors;
if (gmail.connected && !gmailImap.connected) {
return connectors.filter((c) => c.connector_id !== 'gmail_imap');
}
if (gmailImap.connected && !gmail.connected) {
return connectors.filter((c) => c.connector_id !== 'gmail');
}
if (gmail.connected && gmailImap.connected) {
const dropId = gmail.chunks >= gmailImap.chunks ? 'gmail_imap' : 'gmail';
return connectors.filter((c) => c.connector_id !== dropId);
}
// Neither connected — show only the IMAP card as the default flow.
return connectors.filter((c) => c.connector_id !== 'gmail');
})();
const connected = unifiedConnectors.filter((c) => c.connected);
const notConnectedBase = unifiedConnectors.filter((c) => !c.connected);
// Always show the upload card in the not-connected list (it has no backend connector)
const uploadEntry = { connector_id: 'upload', display_name: 'Upload / Paste', connected: false, chunks: 0 };
const notConnected = notConnectedBase.some((c) => c.connector_id === 'upload')
@@ -642,10 +786,9 @@ function DataSourcesSection() {
</div>
<div className="flex flex-col gap-2">
{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 (
<div
@@ -663,7 +806,7 @@ function DataSourcesSection() {
}}>
<div style={{ flex: 1, minWidth: 0 }}>
<div className="font-semibold" style={{ fontSize: 14, fontWeight: 600, color: 'var(--color-text)' }}>
{c.display_name}
{meta?.display_name ?? c.display_name}
</div>
<SyncStatusDisplay
chunks={c.chunks}
@@ -674,60 +817,23 @@ function DataSourcesSection() {
/>
</div>
<button
onClick={() => setExpandedId(isReconnecting ? null : c.connector_id)}
onClick={() => handleDisconnect(c.connector_id)}
disabled={disconnectingId === c.connector_id}
className="hud-label"
style={{
padding: '6px 12px',
background: 'transparent',
color: 'var(--color-text-secondary)',
border: '1px solid var(--color-border)',
borderRadius: 4, cursor: 'pointer',
borderRadius: 4,
cursor: disconnectingId === c.connector_id ? 'default' : 'pointer',
letterSpacing: '0.15em',
opacity: disconnectingId === c.connector_id ? 0.5 : 1,
}}
>
{isReconnecting ? 'Cancel' : 'Reconnect'}
{disconnectingId === c.connector_id ? 'Disconnecting…' : 'Disconnect'}
</button>
</div>
{isReconnecting && meta?.steps && (
<div style={{ borderTop: '1px solid var(--color-border)', padding: 12 }}>
<div style={{ fontSize: 12, color: 'var(--color-warning)', marginBottom: 8 }}>
Re-enter credentials to reconnect this source.
</div>
{meta.steps.map((step, i) => (
<div
key={i}
style={{
background: 'var(--color-bg)',
border: '1px solid var(--color-border)',
borderRadius: 6, padding: 10,
marginBottom: 8,
}}
>
<div style={{ color: 'var(--color-accent-purple)', fontSize: 10, fontWeight: 600, marginBottom: 3 }}>
STEP {i + 1}
</div>
<div style={{ fontSize: 12, marginBottom: step.url ? 4 : 0 }}>{step.label}</div>
{step.url && (
<a
href={step.url}
target="_blank"
rel="noopener noreferrer"
style={{ color: 'var(--color-accent)', fontSize: 11, textDecoration: 'underline' }}
>
{step.urlLabel || 'Open'} &rarr;
</a>
)}
</div>
))}
{meta.inputFields && (
<InlineConnectForm
fields={meta.inputFields}
loading={loading}
onSubmit={(req) => handleConnect(c.connector_id, req)}
/>
)}
</div>
)}
</div>
);
})}
@@ -744,7 +850,7 @@ function DataSourcesSection() {
</div>
<div className="grid grid-cols-2 gap-2">
{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() {
>
<div style={{ flex: 1, minWidth: 0 }}>
<div className="font-semibold" style={{ fontSize: 14, fontWeight: 600, color: 'var(--color-text)' }}>
{c.display_name}
{meta?.display_name ?? c.display_name}
</div>
<div style={{ fontSize: 11, color: 'var(--color-text-tertiary)', marginTop: 2 }}>
Not connected
{meta?.description ?? 'Not connected'}
</div>
</div>
<span style={{ color: 'var(--color-text-secondary)', fontSize: 12, fontWeight: 500 }}>
@@ -822,6 +928,12 @@ function DataSourcesSection() {
onSubmit={(req) => handleConnect(c.connector_id, req)}
/>
)}
{c.connector_id === 'gmail_imap' && (
<GmailOAuthAdvanced
loading={loading && connectingId === 'gmail'}
onConnect={(req) => handleConnect('gmail', req)}
/>
)}
{meta?.troubleshooting && (
<details className="mt-2">
<summary className="text-[11px] cursor-pointer" style={{ color: 'var(--color-text-tertiary)' }}>
+12 -11
View File
@@ -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' },
+57
View File
@@ -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 };
+2 -1
View File
@@ -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',
},
},
});
+5
View File
@@ -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 = [
+20 -6
View File
@@ -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:
+30
View File
@@ -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()
+683
View File
@@ -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": "<ISO 8601>", "end": "<ISO 8601>"}}`. 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:<hex_message_id>"`` 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/<id>`` permalink.
- **RFC822 Message-ID** (``<CABCD@mail.gmail.com>``) — 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",
]
+14 -1
View File
@@ -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
+8 -2
View File
@@ -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)
+8 -1
View File
@@ -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:
+215
View File
@@ -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
+9
View File
@@ -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)
+229 -115
View File
@@ -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)))
+156
View File
@@ -0,0 +1,156 @@
"""Dense embedding clients for the IngestionPipeline.
A thin HTTP wrapper around a local `Ollama <https://ollama.com>`_ 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",
]
+11 -7
View File
@@ -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,
+5 -4
View File
@@ -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:
+11 -5
View File
@@ -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:
+161 -21
View File
@@ -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 <script>, <style>, and <head> contents entirely so CSS rules and
JS payloads don't pollute the text. Inserts a newline at each block-level
tag boundary so the downstream chunker still has paragraph-ish breaks
to split on; without this, a single <div>-wrapped marketing email
becomes one giant unsplittable chunk.
"""
_SKIP_TAGS = {"script", "style", "head", "title", "meta", "link"}
_BLOCK_TAGS = {
"p", "div", "br", "li", "ul", "ol", "tr", "td", "table",
"h1", "h2", "h3", "h4", "h5", "h6", "blockquote", "hr",
"article", "section", "header", "footer", "pre",
}
def __init__(self) -> None:
super().__init__(convert_charrefs=True)
self._parts: List[str] = []
self._skip_depth = 0
def handle_starttag(
self, tag: str, attrs: List[Tuple[str, Optional[str]]]
) -> None:
if tag in self._SKIP_TAGS:
self._skip_depth += 1
elif tag in self._BLOCK_TAGS and self._skip_depth == 0:
self._parts.append("\n")
def handle_endtag(self, tag: str) -> None:
if tag in self._SKIP_TAGS:
self._skip_depth = max(0, self._skip_depth - 1)
def handle_data(self, data: str) -> None:
if self._skip_depth == 0:
self._parts.append(data)
def get_text(self) -> str:
text = "".join(self._parts)
# Collapse runs of horizontal whitespace and excess blank lines so
# the chunker sees clean paragraphs rather than walls of \n.
text = re.sub(r"[ \t\r\f\v]+", " ", text)
text = re.sub(r"\n[ \t]*", "\n", text)
text = re.sub(r"\n{3,}", "\n\n", text)
return text.strip()
def _html_to_text(html: str) -> str:
"""Convert HTML to plain text. Malformed input returns best-effort output."""
extractor = _HTMLTextExtractor()
try:
extractor.feed(html)
extractor.close()
except Exception: # noqa: BLE001
# Parser exceptions still leave partial output in self._parts.
pass
return extractor.get_text()
def _decode_body(payload: Dict[str, Any]) -> str:
"""Decode the message body from a Gmail payload dict.
Handles both simple payloads (``body.data``) and multipart messages
by recursively searching for a ``text/plain`` part.
Multipart messages prefer text/plain. When only text/html is available
(common for marketing emails), the HTML is stripped to plain text so
downstream chunkers and embeddings don't ingest raw markup.
"""
mime_type: str = payload.get("mimeType", "")
if mime_type.startswith("multipart/"):
# Search parts for text/plain first, then any text/* fallback
parts: List[Dict[str, Any]] = payload.get("parts", [])
# Prefer text/plain when both alternatives are present.
for part in parts:
if part.get("mimeType", "").startswith("text/plain"):
return _decode_body(part)
# Fallback: recurse into first part
# Fall back to text/html (which gets stripped at the leaf branch).
for part in parts:
if part.get("mimeType", "").startswith("text/html"):
return _decode_body(part)
# Last resort: recurse into the first part regardless of type.
if parts:
return _decode_body(parts[0])
return ""
@@ -144,10 +227,14 @@ def _decode_body(payload: Dict[str, Any]) -> str:
# Gmail uses URL-safe base64 without padding
padded = body_data + "=" * (-len(body_data) % 4)
try:
return base64.urlsafe_b64decode(padded).decode("utf-8", errors="replace")
decoded = base64.urlsafe_b64decode(padded).decode("utf-8", errors="replace")
except Exception: # noqa: BLE001
return ""
if mime_type.startswith("text/html"):
return _html_to_text(decoded)
return decoded
def _parse_date(date_str: str) -> datetime:
"""Parse an RFC 2822 email date string into a :class:`~datetime.datetime`.
@@ -162,6 +249,31 @@ def _parse_date(date_str: str) -> datetime:
return datetime.now()
def _normalize_addresses(raw: str) -> List[str]:
"""Extract lowercase email addresses from a comma-separated header value.
Uses :func:`email.utils.getaddresses` so multi-recipient ``To``/``Cc``
headers are handled correctly. Addresses that fail to parse are dropped.
"""
if not raw:
return []
return [addr.lower() for _, addr in email.utils.getaddresses([raw]) if addr]
# Order matters: a message tagged both SENT and INBOX (rare) reads as SENT,
# the more specific origin. INBOX is last so it acts as the default.
_PRIMARY_LABELS = ("SENT", "DRAFT", "SPAM", "TRASH", "INBOX")
def _select_channel(label_ids: List[str]) -> Optional[str]:
"""Pick the most specific Gmail system folder this message lives in."""
label_set = set(label_ids)
for label in _PRIMARY_LABELS:
if label in label_set:
return label
return None
# ---------------------------------------------------------------------------
# GmailConnector
# ---------------------------------------------------------------------------
@@ -245,26 +357,34 @@ class GmailConnector(BaseConnector):
cursor:
``nextPageToken`` from a previous sync to resume pagination.
"""
# Existence check only — the actual access token is reloaded on every
# API call by _call_with_refresh so a mid-sync refresh is picked up
# transparently.
tokens = load_tokens(self._credentials_path)
if not tokens:
return
token: str = tokens.get("token", tokens.get("access_token", ""))
if not token:
if not (tokens.get("access_token") or tokens.get("token")):
return
query = "category:primary"
# Default to no filter so SENT, labeled, and category-tabbed mail
# all flow in. The previous "category:primary" default excluded
# ~95% of a typical mailbox (sent mail, Promotions, Updates, etc.)
# which made any C2-style "what did I say to X" query impossible.
query_parts: List[str] = []
if since is not None:
# Gmail's after: operator accepts Unix epoch seconds.
epoch = int(since.timestamp())
query = f"category:primary after:{epoch}"
query_parts.append(f"after:{int(since.timestamp())}")
query = " ".join(query_parts)
page_token: Optional[str] = cursor
synced = 0
while True:
list_resp = _gmail_api_list_messages(
token, page_token=page_token, query=query
list_resp = _call_with_refresh(
_gmail_api_list_messages,
self._credentials_path,
page_token=page_token,
query=query,
)
messages: List[Dict[str, Any]] = list_resp.get("messages", [])
@@ -273,39 +393,59 @@ class GmailConnector(BaseConnector):
if not msg_id:
continue
msg = _gmail_api_get_message(token, msg_id)
msg = _call_with_refresh(
_gmail_api_get_message,
self._credentials_path,
msg_id,
)
payload: Dict[str, Any] = msg.get("payload", {})
headers: List[Dict[str, str]] = payload.get("headers", [])
from_header = _extract_header(headers, "From")
to_header = _extract_header(headers, "To")
cc_header = _extract_header(headers, "Cc")
subject = _extract_header(headers, "Subject")
date_str = _extract_header(headers, "Date")
to_header = _extract_header(headers, "To")
rfc_message_id = _extract_header(headers, "Message-ID")
body = _decode_body(payload)
timestamp = _parse_date(date_str)
# Raw header values, exactly as Gmail returned them — preserved
# so re-normalisation against an updated alias map doesn't need
# a re-fetch from the API.
participants_raw: List[str] = [
h for h in (from_header, to_header, cc_header) if h
]
# Lowercase email addresses, multi-recipient-aware.
participants: List[str] = []
if from_header:
participants.append(from_header)
if to_header:
participants.append(to_header)
for header in (from_header, to_header, cc_header):
participants.extend(_normalize_addresses(header))
label_ids: List[str] = msg.get("labelIds", [])
channel = _select_channel(label_ids)
thread_id: Optional[str] = msg.get("threadId")
doc = Document(
doc_id=f"gmail:{msg_id}",
source="gmail",
source_id=msg_id,
doc_type="email",
content=body,
title=subject,
author=from_header,
participants=participants,
participants_raw=participants_raw,
timestamp=timestamp,
thread_id=thread_id,
channel=channel,
metadata={
"message_id": msg_id,
"labels": msg.get("labelIds", []),
"rfc_message_id": rfc_message_id,
"labels": label_ids,
"snippet": msg.get("snippet", ""),
"history_id": msg.get("historyId", ""),
"size_estimate": msg.get("sizeEstimate", 0),
},
)
synced += 1
+19 -11
View File
@@ -92,12 +92,14 @@ class GmailIMAPConnector(BaseConnector):
credentials_path: str = "",
*,
imap_host: str = "",
max_messages: int = 5000,
max_messages: Optional[int] = None,
) -> None:
self._email = email_address
self._password = app_password
self._credentials_path = credentials_path or _DEFAULT_CREDENTIALS_PATH
self._imap_host = imap_host or self._default_imap_host
# ``None`` means "no cap" — the full inbox is indexed. A positive
# value is still honored for tests that want a bounded scan.
self._max_messages = max_messages
self._items_synced = 0
self._items_total = 0
@@ -156,21 +158,27 @@ class GmailIMAPConnector(BaseConnector):
imap.select("INBOX", readonly=True)
# Build search criteria
if since:
date_str = since.strftime("%d-%b-%Y")
_, data = imap.search(None, f"SINCE {date_str}")
else:
_, data = imap.search(None, "ALL")
# Always SEARCH ALL. IMAP has no native cursor that survives a
# server restart, so applying the SyncEngine's ``since`` filter
# during a partial backfill would silently skip the older
# unprocessed messages. The pipeline-level dedup (_seen_doc_ids
# set + INSERT OR IGNORE in KnowledgeStore) makes re-scanning
# already-indexed messages cheap, so resume is correct as long
# as we keep enumerating the full inbox.
_, data = imap.search(None, "ALL")
msg_ids = data[0].split()
self._items_total = len(msg_ids)
# Take the most recent N messages
recent = msg_ids[-self._max_messages :]
# Newest-first iteration so Deep Research becomes useful while
# the long tail of older mail finishes indexing in the
# background. IMAP returns sequence numbers in arrival order
# (oldest -> newest), so reversing puts the most recent first.
ordered = list(reversed(msg_ids))
if self._max_messages is not None and self._max_messages > 0:
ordered = ordered[: self._max_messages]
synced = 0
for mid in recent:
for mid in ordered:
try:
_, msg_data = imap.fetch(mid, "(RFC822)")
raw = msg_data[0][1]
+130
View File
@@ -0,0 +1,130 @@
"""Shared Google OAuth helpers: access token read + one-shot 401 refresh.
All Google connectors (Gmail, Calendar, Contacts, Drive, Tasks) authenticate
with the same OAuth flow and store identical token payloads at
``~/.openjarvis/connectors/*.json`` — typically a shared ``google.json`` file
plus per-product copies. They all need the same refresh-on-401 behavior, so
the wrapper lives here instead of being duplicated per connector.
Use ``call_with_refresh(api_fn, credentials_path, *args, **kwargs)`` around
any token-taking API helper. On a 401 the wrapper exchanges the stored
``refresh_token`` for a new ``access_token``, updates the credentials file,
and retries the call once. All other status codes propagate.
"""
from __future__ import annotations
import logging
from typing import Any, Callable, Dict
import httpx
from openjarvis.connectors.oauth import load_tokens, save_tokens
logger = logging.getLogger(__name__)
_GOOGLE_TOKEN_ENDPOINT = "https://oauth2.googleapis.com/token"
class GoogleAuthError(RuntimeError):
"""Raised when Google credentials are missing or refresh-token grant fails."""
def current_access_token(credentials_path: str) -> str:
"""Return the current access token from the credentials file (empty if absent)."""
tokens = load_tokens(credentials_path) or {}
return tokens.get("access_token", tokens.get("token", ""))
def refresh_access_token(credentials_path: str) -> str:
"""Exchange the stored refresh_token for a fresh access_token and persist it.
Returns the new access_token. Raises :class:`GoogleAuthError` when the
credentials file is missing, lacks a refresh_token / client credentials,
or when Google rejects the refresh grant (e.g. the refresh_token has been
revoked and the user needs to re-authenticate).
"""
tokens = load_tokens(credentials_path)
if not tokens:
raise GoogleAuthError(
f"No credentials at {credentials_path}; re-run the connector OAuth flow."
)
refresh_token = tokens.get("refresh_token", "")
client_id = tokens.get("client_id", "")
client_secret = tokens.get("client_secret", "")
if not (refresh_token and client_id and client_secret):
raise GoogleAuthError(
"Stored Google credentials are missing refresh_token / client_id / "
"client_secret; re-run the connector OAuth flow to mint a full token."
)
resp = httpx.post(
_GOOGLE_TOKEN_ENDPOINT,
data={
"client_id": client_id,
"client_secret": client_secret,
"refresh_token": refresh_token,
"grant_type": "refresh_token",
},
timeout=30.0,
)
if resp.status_code != 200:
raise GoogleAuthError(
f"Google token refresh failed ({resp.status_code}): {resp.text[:200]}"
)
payload = resp.json()
new_token = payload.get("access_token", "")
if not new_token:
raise GoogleAuthError(
"Google token refresh returned 200 but no access_token in payload."
)
tokens["access_token"] = new_token
# Keep the legacy "token" key in sync for older code paths that read it.
tokens["token"] = new_token
if "expires_in" in payload:
tokens["expires_in"] = payload["expires_in"]
save_tokens(credentials_path, tokens)
logger.info(
"Refreshed Google access token (expires_in=%s)", payload.get("expires_in")
)
return new_token
def call_with_refresh(
api_fn: Callable[..., Dict[str, Any]],
credentials_path: str,
*args: Any,
**kwargs: Any,
) -> Dict[str, Any]:
"""Invoke ``api_fn(token, *args, **kwargs)`` with one-shot 401 auto-refresh.
Loads the current access token from disk, calls the helper, and if Google
returns 401 (the access token has expired or been revoked) uses the stored
refresh_token to mint a new access_token, updates the credentials file,
and retries the call exactly once.
Any other ``HTTPStatusError`` is re-raised unchanged — auth-related retries
end here; transient 5xx / timeout retries belong further up the stack.
"""
token = current_access_token(credentials_path)
try:
return api_fn(token, *args, **kwargs)
except httpx.HTTPStatusError as exc:
if exc.response is None or exc.response.status_code != 401:
raise
logger.info(
"Google returned 401 on %s — refreshing access token and retrying.",
getattr(api_fn, "__name__", "<api_fn>"),
)
new_token = refresh_access_token(credentials_path)
return api_fn(new_token, *args, **kwargs)
__all__ = [
"GoogleAuthError",
"current_access_token",
"refresh_access_token",
"call_with_refresh",
]
+11 -5
View File
@@ -12,6 +12,7 @@ from typing import Any, Dict, Iterator, 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 load_tokens, resolve_google_credentials
from openjarvis.core.config import DEFAULT_CONFIG_DIR
from openjarvis.core.registry import ConnectorRegistry
@@ -62,10 +63,10 @@ class GoogleTasksConnector(BaseConnector):
def sync(
self, *, since: Optional[datetime] = None, cursor: Optional[str] = None
) -> Iterator[Document]:
token = self._get_access_token()
# List all task lists first
task_lists = _tasks_api_get(token, "users/@me/lists")
# call_with_refresh handles the access-token read + one-shot 401 retry.
task_lists = call_with_refresh(
_tasks_api_get, str(self._credentials_path), "users/@me/lists"
)
for tl in task_lists.get("items", []):
tl_id = tl["id"]
@@ -79,7 +80,12 @@ class GoogleTasksConnector(BaseConnector):
if since:
params["updatedMin"] = since.isoformat() + "Z"
tasks = _tasks_api_get(token, f"lists/{tl_id}/tasks", params=params)
tasks = call_with_refresh(
_tasks_api_get,
str(self._credentials_path),
f"lists/{tl_id}/tasks",
params=params,
)
for task in tasks.get("items", []):
due = task.get("due", "")
+471
View File
@@ -0,0 +1,471 @@
"""Hybrid retrieval over the KnowledgeStore: metadata filter + BM25 + vector cosine.
A single ``search`` entrypoint that the agentic research loop calls as a tool.
Structured WHERE-clause filters (person, time range, sources) narrow the
candidate set, then BM25 (FTS5) and dense cosine similarity score the
survivors. The two ranks are fused with Reciprocal Rank Fusion, which is
robust to the very different score scales the two signals produce
(BM25 ~ [0, 20], cosine ~ [0.4, 0.9] for nomic-embed-text).
Each result is enriched with its thread context: when a hit belongs to a
``thread_id``, the surrounding chunks are attached so the synthesis model
sees the conversation, not an isolated fragment.
Brute-force vector scan is fine at the current corpus size (~5k chunks ×
768 dims fits in ~15 MB and matmuls in <50 ms). Swap in an ANN index when
that stops being true.
"""
from __future__ import annotations
import json
import logging
from dataclasses import dataclass, field
from datetime import datetime
from typing import Any, Dict, List, Optional, Sequence, Tuple
import numpy as np
from openjarvis.connectors.embeddings import OllamaEmbedder, decode_embedding
from openjarvis.connectors.store import KnowledgeStore
logger = logging.getLogger(__name__)
# ---------------------------------------------------------------------------
# Result types
# ---------------------------------------------------------------------------
@dataclass
class SearchHit:
"""A single hybrid-search result with enough context for citation."""
chunk_id: str
document_id: str
chunk_idx: int
title: str
content_snippet: str
source: str
timestamp: str
participants: List[str]
score: float
bm25_score: float
vector_score: float
thread_id: str = ""
thread_context: List[Dict[str, Any]] = field(default_factory=list)
def to_dict(self) -> Dict[str, Any]:
return {
"title": self.title,
"content_snippet": self.content_snippet,
"source": self.source,
"timestamp": self.timestamp,
"participants": self.participants,
"score": round(self.score, 4),
"document_id": self.document_id,
"chunk_idx": self.chunk_idx,
"thread_id": self.thread_id,
"thread_context": self.thread_context,
}
# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------
def _iso(ts: Optional[datetime | str]) -> Optional[str]:
if ts is None:
return None
if isinstance(ts, datetime):
return ts.isoformat()
return str(ts)
def _quote_fts(query: str) -> str:
"""Make a plain user query safe for FTS5 MATCH.
FTS5 treats characters like ``-``, ``:``, ``"`` as operators; the simplest
way to avoid syntax errors on arbitrary user input is to quote each
whitespace-delimited token and OR them together.
"""
tokens = [t for t in query.split() if t]
if not tokens:
return ""
return " OR ".join(f'"{t.replace(chr(34), "")}"' for t in tokens)
def _parse_participants(raw: Any) -> List[str]:
if not raw:
return []
if isinstance(raw, list):
return [str(x) for x in raw]
try:
parsed = json.loads(raw)
return [str(x) for x in parsed] if isinstance(parsed, list) else []
except (json.JSONDecodeError, TypeError):
return []
def _snippet(content: str, max_chars: int = 500) -> str:
flat = content.strip()
if len(flat) <= max_chars:
return flat
return flat[:max_chars].rstrip() + ""
# ---------------------------------------------------------------------------
# HybridSearch
# ---------------------------------------------------------------------------
class HybridSearch:
"""Hybrid BM25 + dense-cosine retrieval over a ``KnowledgeStore``.
Parameters
----------
store:
The store to query.
embedder:
Embedding client used to encode the query. When ``None``, search
falls back to BM25 only and reports ``vector_score=0``.
bm25_weight, vector_weight:
Weights on the two RRF terms. Defaults to 0.5 / 0.5; raise either to
bias retrieval toward lexical or semantic matches.
rrf_k:
RRF damping constant. Larger values flatten the contribution of
deeper ranks; 60 is the canonical value from the original paper.
recall_k:
How deep each individual ranker recalls before fusion. Should be at
least a few times ``limit`` so the fuser has overlap to work with.
"""
def __init__(
self,
store: KnowledgeStore,
embedder: Optional[OllamaEmbedder] = None,
*,
bm25_weight: float = 0.5,
vector_weight: float = 0.5,
rrf_k: int = 60,
recall_k: int = 200,
thread_context_cap: int = 20,
) -> None:
self._store = store
self._embedder = embedder
self._bm25_weight = float(bm25_weight)
self._vector_weight = float(vector_weight)
self._rrf_k = int(rrf_k)
self._recall_k = int(recall_k)
self._thread_context_cap = int(thread_context_cap)
# ------------------------------------------------------------------
# Filter SQL construction
# ------------------------------------------------------------------
def _build_filters(
self,
*,
person: Optional[str],
time_range: Optional[Tuple[Optional[datetime], Optional[datetime]]],
sources: Optional[Sequence[str]],
alias: str = "",
) -> Tuple[str, List[Any]]:
"""Return ``(where_fragment, params)`` for the structured filters.
``person`` is matched against the participants_raw JSON via LIKE so a
substring of a name or email address is enough — handy when the user
says "Kelly" rather than "kelly@example.com".
``alias`` qualifies every column reference (e.g. ``kc.`` when joining
against the FTS virtual table which also has ``author`` and ``title``
columns and would otherwise produce an "ambiguous column" error).
"""
prefix = f"{alias}." if alias else ""
clauses: List[str] = [f"{prefix}deleted_at IS NULL"]
params: List[Any] = []
if person:
clauses.append(
f"({prefix}participants_raw LIKE ? OR {prefix}participants LIKE ? "
f"OR {prefix}author LIKE ?)"
)
needle = f"%{person}%"
params.extend([needle, needle, needle])
if time_range:
start, end = time_range
if start is not None:
clauses.append(f"{prefix}timestamp >= ?")
params.append(_iso(start))
if end is not None:
clauses.append(f"{prefix}timestamp <= ?")
params.append(_iso(end))
if sources:
placeholders = ",".join("?" for _ in sources)
clauses.append(f"{prefix}source IN ({placeholders})")
params.extend(sources)
return " AND ".join(clauses), params
# ------------------------------------------------------------------
# BM25 leg
# ------------------------------------------------------------------
def _bm25_recall(
self, query: str, filter_sql: str, filter_params: List[Any]
) -> List[Tuple[str, float]]:
"""Return ``[(chunk_id, bm25_score), ...]`` from FTS5."""
fts_query = _quote_fts(query)
if not fts_query:
return []
sql = f"""
SELECT kc.id, abs(bm25(knowledge_fts)) AS score
FROM knowledge_fts
JOIN knowledge_chunks kc ON knowledge_fts.rowid = kc.rowid
WHERE knowledge_fts MATCH ?
AND {filter_sql}
ORDER BY score DESC
LIMIT ?
"""
try:
rows = self._store._conn.execute(
sql, [fts_query, *filter_params, self._recall_k]
).fetchall()
except Exception as exc: # noqa: BLE001
logger.warning("hybrid_search: BM25 leg failed (%s)", exc)
return []
return [(row["id"], float(row["score"])) for row in rows]
# ------------------------------------------------------------------
# Vector leg
# ------------------------------------------------------------------
def _vector_recall(
self, query: str, filter_sql: str, filter_params: List[Any]
) -> List[Tuple[str, float]]:
"""Return ``[(chunk_id, cosine_score), ...]`` from a brute-force scan."""
if self._embedder is None:
return []
q_blob = self._embedder.embed(query)
q_vec = decode_embedding(q_blob)
if q_vec is None or q_vec.size == 0:
return []
q_norm = float(np.linalg.norm(q_vec))
if q_norm == 0.0:
return []
q_unit = q_vec / q_norm
sql = f"""
SELECT id, embedding
FROM knowledge_chunks
WHERE embedding IS NOT NULL AND {filter_sql}
"""
rows = self._store._conn.execute(sql, filter_params).fetchall()
if not rows:
return []
ids: List[str] = []
vecs: List[np.ndarray] = []
for row in rows:
vec = decode_embedding(row["embedding"])
if vec is None or vec.size != q_unit.size:
continue
ids.append(row["id"])
vecs.append(vec)
if not ids:
return []
mat = np.vstack(vecs).astype(np.float32, copy=False)
norms = np.linalg.norm(mat, axis=1)
norms[norms == 0.0] = 1.0
mat = mat / norms[:, None]
scores = mat @ q_unit
# Top-recall_k
if len(scores) > self._recall_k:
top_idx = np.argpartition(-scores, self._recall_k)[: self._recall_k]
top_idx = top_idx[np.argsort(-scores[top_idx])]
else:
top_idx = np.argsort(-scores)
return [(ids[int(i)], float(scores[int(i)])) for i in top_idx]
# ------------------------------------------------------------------
# Fusion
# ------------------------------------------------------------------
def _fuse(
self,
bm25: List[Tuple[str, float]],
vector: List[Tuple[str, float]],
) -> List[Tuple[str, float, float, float]]:
"""Reciprocal Rank Fusion across the two rankers.
Returns ``[(chunk_id, fused, bm25_score, vector_score), ...]``
sorted by ``fused`` descending.
"""
bm25_rank = {cid: i + 1 for i, (cid, _) in enumerate(bm25)}
vec_rank = {cid: i + 1 for i, (cid, _) in enumerate(vector)}
bm25_scores = {cid: s for cid, s in bm25}
vec_scores = {cid: s for cid, s in vector}
candidates = set(bm25_rank) | set(vec_rank)
out: List[Tuple[str, float, float, float]] = []
for cid in candidates:
fused = 0.0
if cid in bm25_rank:
fused += self._bm25_weight / (self._rrf_k + bm25_rank[cid])
if cid in vec_rank:
fused += self._vector_weight / (self._rrf_k + vec_rank[cid])
out.append(
(cid, fused, bm25_scores.get(cid, 0.0), vec_scores.get(cid, 0.0))
)
out.sort(key=lambda r: -r[1])
return out
# ------------------------------------------------------------------
# Thread enrichment
# ------------------------------------------------------------------
def _thread_context(
self, thread_id: str, anchor_chunk_id: str
) -> List[Dict[str, Any]]:
"""Fetch sibling chunks for ``thread_id`` (capped at ``thread_context_cap``).
When the thread is longer than the cap, return a centred window around
the anchor so the most relevant chunk is always present.
"""
if not thread_id:
return []
rows = self._store._conn.execute(
"""
SELECT id, chunk_index, content, timestamp, author
FROM knowledge_chunks
WHERE thread_id = ? AND deleted_at IS NULL
ORDER BY timestamp ASC, chunk_index ASC
""",
(thread_id,),
).fetchall()
if not rows:
return []
cap = self._thread_context_cap
if len(rows) > cap:
anchor_idx = next(
(i for i, r in enumerate(rows) if r["id"] == anchor_chunk_id),
len(rows) // 2,
)
half = cap // 2
lo = max(0, anchor_idx - half)
hi = min(len(rows), lo + cap)
lo = max(0, hi - cap)
rows = rows[lo:hi]
return [
{
"chunk_idx": int(r["chunk_index"]),
"timestamp": r["timestamp"] or "",
"author": r["author"] or "",
"snippet": _snippet(r["content"], 240),
}
for r in rows
]
# ------------------------------------------------------------------
# Public entry point
# ------------------------------------------------------------------
def search(
self,
query: str,
*,
person: Optional[str] = None,
time_range: Optional[Tuple[Optional[datetime], Optional[datetime]]] = None,
sources: Optional[Sequence[str]] = None,
limit: int = 20,
) -> List[SearchHit]:
"""Run the hybrid pipeline and return up to ``limit`` hits.
See module docstring for ranking semantics. ``query`` may be empty
when callers want a pure metadata filter (e.g. "all mail from X in
May") — in that case only the vector leg runs (and only if an
embedder is configured); if neither leg yields anything the
structured filter is applied directly and the most recent rows are
returned.
"""
bm25_filter_sql, bm25_filter_params = self._build_filters(
person=person, time_range=time_range, sources=sources, alias="kc"
)
unaliased_filter_sql, unaliased_filter_params = self._build_filters(
person=person, time_range=time_range, sources=sources
)
bm25 = (
self._bm25_recall(query, bm25_filter_sql, bm25_filter_params)
if query.strip()
else []
)
vector = (
self._vector_recall(query, unaliased_filter_sql, unaliased_filter_params)
if query.strip()
else []
)
fused = self._fuse(bm25, vector)
# Metadata-only fallback: empty query, or both legs produced nothing
# despite a non-empty query. Return the most recent rows matching the
# filter so the agent still gets a useful corpus snapshot.
if not fused:
sql = f"""
SELECT id FROM knowledge_chunks
WHERE {unaliased_filter_sql}
ORDER BY timestamp DESC, created_at DESC
LIMIT ?
"""
rows = self._store._conn.execute(
sql, [*unaliased_filter_params, limit]
).fetchall()
fused = [(row["id"], 0.0, 0.0, 0.0) for row in rows]
# Materialise the top-N rows in one IN-clause round trip.
top = fused[:limit]
if not top:
return []
ids = [cid for cid, *_ in top]
placeholders = ",".join("?" for _ in ids)
meta_rows = self._store._conn.execute(
f"""
SELECT id, doc_id, content, source, title, author, participants,
timestamp, thread_id, chunk_index
FROM knowledge_chunks
WHERE id IN ({placeholders})
""",
ids,
).fetchall()
by_id = {r["id"]: r for r in meta_rows}
hits: List[SearchHit] = []
for chunk_id, fused_score, bm25_score, vec_score in top:
r = by_id.get(chunk_id)
if r is None:
continue
hits.append(
SearchHit(
chunk_id=chunk_id,
document_id=r["doc_id"],
chunk_idx=int(r["chunk_index"]),
title=r["title"] or "",
content_snippet=_snippet(r["content"]),
source=r["source"] or "",
timestamp=r["timestamp"] or "",
participants=_parse_participants(r["participants"]),
score=fused_score,
bm25_score=bm25_score,
vector_score=vec_score,
thread_id=r["thread_id"] or "",
thread_context=self._thread_context(r["thread_id"] or "", chunk_id),
)
)
return hits
__all__ = ["HybridSearch", "SearchHit"]
+92 -3
View File
@@ -13,12 +13,49 @@ Typical usage::
from __future__ import annotations
import hashlib
import time
from typing import TYPE_CHECKING, Iterable, Optional
from openjarvis.connectors._stubs import Attachment, Document
from openjarvis.connectors.chunker import SemanticChunker
from openjarvis.connectors.embeddings import OllamaEmbedder
from openjarvis.connectors.store import KnowledgeStore
def _namespace_thread_id(source: str, thread_id: Optional[str]) -> Optional[str]:
"""Prefix ``thread_id`` with ``{source}:`` so it can't collide across sources.
Idempotent: if the input already starts with ``{source}:`` it is returned
unchanged. Centralised here (rather than per-connector) so a new connector
author can't forget to namespace.
"""
if not thread_id:
return None
prefix = f"{source}:"
if thread_id.startswith(prefix):
return thread_id
return f"{prefix}{thread_id}"
def _derive_source_id(doc: Document) -> str:
"""Return the connector-set ``source_id`` or extract it from ``doc_id``.
Many existing connectors compose ``doc_id = f"{source}:{native_id}"``;
this strips the prefix so storage can index the native ID directly.
"""
if doc.source_id:
return doc.source_id
prefix = f"{doc.source}:"
if doc.doc_id.startswith(prefix):
return doc.doc_id[len(prefix):]
return doc.doc_id
def _content_hash(text: str) -> str:
"""SHA-256 hex digest of UTF-8-encoded chunk content."""
return hashlib.sha256(text.encode("utf-8")).hexdigest()
if TYPE_CHECKING:
from openjarvis.connectors.attachment_store import AttachmentStore
@@ -36,6 +73,12 @@ class IngestionPipeline:
Optional ``AttachmentStore`` for persisting attachment blobs and
extracting text from supported MIME types (PDF, plain text, etc.).
When ``None`` (default) attachments are silently ignored.
embedder:
Optional embedding client (e.g. ``OllamaEmbedder``). When provided,
every chunk is embedded at ingest time and the resulting float32
vector is written to the ``embedding`` BLOB column alongside
``embedding_model_version``. ``None`` (default) skips embedding so
in-memory tests and offline runs don't depend on a sidecar daemon.
"""
def __init__(
@@ -44,10 +87,12 @@ class IngestionPipeline:
*,
max_tokens: int = 512,
attachment_store: Optional[AttachmentStore] = None,
embedder: Optional[OllamaEmbedder] = None,
) -> None:
self._store = store
self._chunker = SemanticChunker(max_tokens=max_tokens)
self._attachment_store = attachment_store
self._embedder = embedder
self._seen_doc_ids: set[str] = set()
self._load_existing_doc_ids()
@@ -62,6 +107,20 @@ class IngestionPipeline:
).fetchall()
self._seen_doc_ids = {r[0] for r in rows}
def _embed_chunk(self, content: str) -> tuple[Optional[bytes], str]:
"""Return ``(embedding_bytes, model_version)`` for a chunk.
Returns ``(None, "")`` when no embedder is configured or the embedder
fails — ingestion continues with the lexical-only row, so a flaky
local daemon never blocks a sync.
"""
if self._embedder is None:
return None, ""
emb = self._embedder.embed(content)
if emb is None:
return None, ""
return emb, self._embedder.model_version
def _extract_attachment_text(self, att: Attachment) -> str:
"""Extract text from an attachment.
@@ -109,15 +168,22 @@ class IngestionPipeline:
if doc.doc_id in self._seen_doc_ids:
continue
# Compute v1 provenance fields once per document.
namespaced_thread = _namespace_thread_id(doc.source, doc.thread_id)
source_id = _derive_source_id(doc)
ingest_epoch = time.time()
# Build the parent metadata dict that will be inherited by every
# chunk produced from this document.
parent_meta = {
"title": doc.title,
"author": doc.author,
"source": doc.source,
"source_id": source_id,
"doc_type": doc.doc_type,
"url": doc.url or "",
"thread_id": doc.thread_id or "",
"thread_id": namespaced_thread or "",
"channel": doc.channel or "",
}
# Merge any extra connector-level metadata (without overwriting
# the standard provenance fields set above).
@@ -137,19 +203,27 @@ class IngestionPipeline:
)
for chunk in chunks:
embedding_bytes, embedding_version = self._embed_chunk(chunk.content)
self._store.store(
content=chunk.content,
source=doc.source,
source_id=source_id,
doc_type=doc.doc_type,
doc_id=doc.doc_id,
title=doc.title,
author=doc.author,
participants=doc.participants,
participants_raw=doc.participants_raw,
timestamp=timestamp_str,
thread_id=doc.thread_id,
thread_id=namespaced_thread,
channel=doc.channel,
url=doc.url,
metadata=chunk.metadata,
chunk_index=chunk.index,
content_hash=_content_hash(chunk.content),
embedding=embedding_bytes,
embedding_model_version=embedding_version,
last_synced=ingest_epoch,
)
chunks_stored += 1
@@ -179,20 +253,35 @@ class IngestionPipeline:
"sha256": sha,
},
)
# Synthetic source_id keeps attachment chunks distinct
# from body chunks under the UNIQUE(source, source_id,
# chunk_index) constraint while still letting them share
# a parent doc_id for dedup and blob linkage.
att_source_id = f"{source_id}#{att.filename}"
for chunk in att_chunks:
embedding_bytes, embedding_version = self._embed_chunk(
chunk.content
)
self._store.store(
content=chunk.content,
source=doc.source,
source_id=att_source_id,
doc_type=doc.doc_type,
doc_id=doc.doc_id,
title=f"{doc.title} [{att.filename}]",
author=doc.author,
participants=doc.participants,
participants_raw=doc.participants_raw,
timestamp=timestamp_str,
thread_id=doc.thread_id,
thread_id=namespaced_thread,
channel=doc.channel,
url=doc.url,
metadata=chunk.metadata,
chunk_index=chunk.index,
content_hash=_content_hash(chunk.content),
embedding=embedding_bytes,
embedding_model_version=embedding_version,
last_synced=ingest_epoch,
)
chunks_stored += 1
+115 -16
View File
@@ -31,19 +31,40 @@ CREATE TABLE IF NOT EXISTS knowledge_chunks (
doc_id TEXT NOT NULL,
content TEXT NOT NULL,
source TEXT NOT NULL DEFAULT '',
source_id TEXT NOT NULL DEFAULT '',
doc_type TEXT NOT NULL DEFAULT '',
title TEXT NOT NULL DEFAULT '',
author TEXT NOT NULL DEFAULT '',
participants TEXT NOT NULL DEFAULT '[]',
participants_raw TEXT NOT NULL DEFAULT '[]',
timestamp TEXT NOT NULL DEFAULT '',
thread_id TEXT NOT NULL DEFAULT '',
channel TEXT NOT NULL DEFAULT '',
url TEXT NOT NULL DEFAULT '',
metadata TEXT NOT NULL DEFAULT '{}',
chunk_index INTEGER NOT NULL DEFAULT 0,
embedding BLOB,
embedding_model_version TEXT NOT NULL DEFAULT '',
content_hash TEXT NOT NULL DEFAULT '',
deleted_at REAL,
last_synced REAL NOT NULL DEFAULT 0,
created_at REAL NOT NULL
);
"""
# Additive migrations: existing installs get new columns via ALTER TABLE.
# Order doesn't matter — each entry is independent and idempotent.
_V1_COLUMNS: List[tuple[str, str]] = [
("source_id", "TEXT NOT NULL DEFAULT ''"),
("participants_raw", "TEXT NOT NULL DEFAULT '[]'"),
("channel", "TEXT NOT NULL DEFAULT ''"),
("embedding", "BLOB"),
("embedding_model_version", "TEXT NOT NULL DEFAULT ''"),
("content_hash", "TEXT NOT NULL DEFAULT ''"),
("deleted_at", "REAL"),
("last_synced", "REAL NOT NULL DEFAULT 0"),
]
_CREATE_FTS_TABLE = """
CREATE VIRTUAL TABLE IF NOT EXISTS knowledge_fts
USING fts5(
@@ -76,12 +97,17 @@ END;
"""
_CREATE_INDEXES = """
CREATE INDEX IF NOT EXISTS idx_kc_source ON knowledge_chunks(source);
CREATE INDEX IF NOT EXISTS idx_kc_doc_type ON knowledge_chunks(doc_type);
CREATE INDEX IF NOT EXISTS idx_kc_author ON knowledge_chunks(author);
CREATE INDEX IF NOT EXISTS idx_kc_timestamp ON knowledge_chunks(timestamp);
CREATE INDEX IF NOT EXISTS idx_kc_thread_id ON knowledge_chunks(thread_id);
CREATE INDEX IF NOT EXISTS idx_kc_doc_id ON knowledge_chunks(doc_id);
CREATE INDEX IF NOT EXISTS idx_kc_source ON knowledge_chunks(source);
CREATE INDEX IF NOT EXISTS idx_kc_doc_type ON knowledge_chunks(doc_type);
CREATE INDEX IF NOT EXISTS idx_kc_author ON knowledge_chunks(author);
CREATE INDEX IF NOT EXISTS idx_kc_timestamp ON knowledge_chunks(timestamp);
CREATE INDEX IF NOT EXISTS idx_kc_thread_id ON knowledge_chunks(thread_id);
CREATE INDEX IF NOT EXISTS idx_kc_doc_id ON knowledge_chunks(doc_id);
CREATE INDEX IF NOT EXISTS idx_kc_source_id ON knowledge_chunks(source_id);
CREATE INDEX IF NOT EXISTS idx_kc_content_hash ON knowledge_chunks(content_hash);
CREATE UNIQUE INDEX IF NOT EXISTS idx_kc_natural_key
ON knowledge_chunks(source, source_id, chunk_index)
WHERE source_id != '';
"""
# ---------------------------------------------------------------------------
@@ -100,6 +126,20 @@ def _to_iso(ts: Optional[Union[datetime, str]]) -> str:
return ts.isoformat()
def _to_epoch(ts: Union[datetime, str, float, int]) -> float:
"""Normalise a timestamp to Unix epoch seconds (float)."""
if isinstance(ts, (int, float)):
return float(ts)
if isinstance(ts, str):
# Treat empty/zero-ish strings as epoch 0; otherwise parse ISO 8601.
if not ts:
return 0.0
return datetime.fromisoformat(ts).timestamp()
if ts.tzinfo is None:
ts = ts.replace(tzinfo=timezone.utc)
return ts.timestamp()
# ---------------------------------------------------------------------------
# KnowledgeStore
# ---------------------------------------------------------------------------
@@ -147,10 +187,27 @@ class KnowledgeStore(MemoryBackend):
self._conn.execute("PRAGMA journal_mode=WAL;")
self._conn.execute("PRAGMA foreign_keys=ON;")
self._conn.executescript(
_CREATE_MAIN_TABLE + _CREATE_FTS_TABLE + _CREATE_TRIGGERS + _CREATE_INDEXES
_CREATE_MAIN_TABLE + _CREATE_FTS_TABLE + _CREATE_TRIGGERS
)
self._migrate_v1_columns()
# Indexes reference the new columns, so they run after migrations.
self._conn.executescript(_CREATE_INDEXES)
self._conn.commit()
def _migrate_v1_columns(self) -> None:
"""Add v1 schema columns to pre-existing tables (idempotent)."""
existing = {
row[1]
for row in self._conn.execute(
"PRAGMA table_info(knowledge_chunks)"
).fetchall()
}
for col_name, col_def in _V1_COLUMNS:
if col_name not in existing:
self._conn.execute(
f"ALTER TABLE knowledge_chunks ADD COLUMN {col_name} {col_def}"
)
# ------------------------------------------------------------------
# MemoryBackend interface
# ------------------------------------------------------------------
@@ -170,6 +227,14 @@ class KnowledgeStore(MemoryBackend):
url: Optional[str] = None,
metadata: Optional[Dict[str, Any]] = None,
chunk_index: int = 0,
# v1 schema additions
source_id: str = "",
participants_raw: Optional[List[str]] = None,
channel: Optional[str] = None,
content_hash: str = "",
embedding: Optional[bytes] = None,
embedding_model_version: str = "",
last_synced: Optional[Union[datetime, str, float]] = None,
) -> str:
"""Persist a content chunk and return its unique chunk id.
@@ -182,50 +247,83 @@ class KnowledgeStore(MemoryBackend):
ts_str = _to_iso(timestamp)
participants_json = json.dumps(participants or [])
participants_raw_json = json.dumps(participants_raw or [])
last_synced_epoch = (
_to_epoch(last_synced) if last_synced is not None else time.time()
)
# Merge provenance fields into metadata for easy access in results
combined_meta: Dict[str, Any] = dict(metadata or {})
combined_meta["chunk_id"] = chunk_id
combined_meta["source"] = source
combined_meta["source_id"] = source_id
combined_meta["doc_type"] = doc_type
combined_meta["doc_id"] = doc_id
combined_meta["title"] = title
combined_meta["author"] = author
combined_meta["participants"] = participants or []
combined_meta["participants_raw"] = participants_raw or []
combined_meta["timestamp"] = ts_str
combined_meta["thread_id"] = thread_id or ""
combined_meta["channel"] = channel or ""
combined_meta["url"] = url or ""
combined_meta["chunk_index"] = chunk_index
combined_meta["content_hash"] = content_hash
combined_meta["embedding_model_version"] = embedding_model_version
meta_json = json.dumps(combined_meta)
self._conn.execute(
cur = self._conn.execute(
"""
INSERT INTO knowledge_chunks
(id, doc_id, content, source, doc_type, title, author,
participants, timestamp, thread_id, url, metadata,
chunk_index, created_at)
VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
INSERT OR IGNORE INTO knowledge_chunks
(id, doc_id, content, source, source_id, doc_type, title, author,
participants, participants_raw, timestamp, thread_id, channel,
url, metadata, chunk_index, embedding, embedding_model_version,
content_hash, last_synced, created_at)
VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
""",
(
chunk_id,
doc_id,
content,
source,
source_id,
doc_type,
title,
author,
participants_json,
participants_raw_json,
ts_str,
thread_id or "",
channel or "",
url or "",
meta_json,
chunk_index,
embedding,
embedding_model_version,
content_hash,
last_synced_epoch,
time.time(),
),
)
self._conn.commit()
# If the natural-key (source, source_id, chunk_index) already existed
# SQLite skipped the insert. Re-runs of the dogfood script or a
# SyncEngine that re-emits a known document hit this path; return
# the existing chunk id so callers don't see two different identities
# for the same row, and suppress the MEMORY_STORE event (nothing new
# was actually written).
if cur.rowcount == 0:
existing = self._conn.execute(
"SELECT id FROM knowledge_chunks "
"WHERE source=? AND source_id=? AND chunk_index=?",
(source, source_id, chunk_index),
).fetchone()
if existing is not None:
return existing["id"]
return chunk_id
get_event_bus().publish(
EventType.MEMORY_STORE,
{
@@ -288,9 +386,10 @@ class KnowledgeStore(MemoryBackend):
filters.append("kc.timestamp <= ?")
params.append(until_str)
where_clause = ""
if filters:
where_clause = "AND " + " AND ".join(filters)
# Always exclude tombstoned rows.
filters.append("kc.deleted_at IS NULL")
where_clause = "AND " + " AND ".join(filters)
# FTS5 bm25() returns negative scores; abs() gives a positive rank
sql = f"""
+8
View File
@@ -192,6 +192,14 @@ class OllamaEngine(InferenceEngine):
"num_ctx": kwargs.get("num_ctx", 8192),
},
}
# Mirror generate()'s default: disable extended thinking unless the
# caller opted in. Qwen3/etc. with thinking on can stall the visible
# stream for 60+ seconds before any tokens reach the client, which
# frontends interpret as a "Load failed" timeout.
if "think" not in kwargs:
payload["think"] = False
elif kwargs["think"] is not None:
payload["think"] = kwargs["think"]
try:
with self._client.stream("POST", "/api/chat", json=payload) as resp:
resp.raise_for_status()
@@ -1421,7 +1421,11 @@ def create_agent_manager_router(
server_config,
)
executor.set_system(system)
executor.execute_tick(agent_id)
# The route handler above already called start_tick() to
# serialize concurrent POSTs; tell the executor not to
# re-acquire, otherwise it bails on its own guard and the
# tick never runs.
executor.execute_tick(agent_id, lock_already_held=True)
except Exception as exc:
logger.error(
"Run-tick failed for agent %s: %s",
+2
View File
@@ -16,6 +16,7 @@ from openjarvis.server.comparison import comparison_router
from openjarvis.server.connectors_router import create_connectors_router
from openjarvis.server.dashboard import dashboard_router
from openjarvis.server.digest_routes import create_digest_router
from openjarvis.server.research_router import router as research_router
from openjarvis.server.routes import router
from openjarvis.server.upload_router import router as upload_router
@@ -287,6 +288,7 @@ def create_app(
app.include_router(create_connectors_router())
app.include_router(create_digest_router())
app.include_router(upload_router)
app.include_router(research_router)
app.include_router(analytics_router)
include_all_routes(app)
+177 -82
View File
@@ -125,6 +125,95 @@ def create_connectors_router():
"chunks": chunks,
}
# ------------------------------------------------------------------
# Background-sync state tracking. Defined here (before the endpoints)
# so that POST /connect can fire-and-forget into the same machinery
# as POST /sync — a single source of truth for "is this connector
# currently syncing", baseline_items, and error translation.
# ------------------------------------------------------------------
_sync_threads: Dict[str, Any] = {}
_sync_state: Dict[str, Dict[str, Any]] = {}
def _translate_sync_error(raw: str) -> str:
"""Map common backend exceptions to a short user-facing message."""
if "401" in raw or "Unauthorized" in raw:
return "Authentication failed — credentials may have expired."
if "403" in raw or "Forbidden" in raw:
return "Permission denied — check API scopes."
if "429" in raw or "Too Many Requests" in raw:
return "Rate limited — wait a minute and try again."
if "timeout" in raw.lower():
return "Connection timed out."
return raw
def _start_sync(connector_id: str, instance: Any) -> str:
"""Spawn a background sync; returns ``"started"`` or ``"already_syncing"``.
Records baseline items in ``_sync_state`` so GET /{id}/sync polls
can compute "new this run" deltas, and translates exceptions into
the same compact error strings the POST /sync handler used to do
inline. Both POST /connect (auto-ingest) and POST /sync (manual)
route through here so they share guard, state, and error handling.
"""
import threading
existing = _sync_threads.get(connector_id)
if existing and existing.is_alive():
return "already_syncing"
# Snapshot the prior checkpoint count so the GET handler can
# report "X new this run" without each client tracking it.
baseline_items = 0
try:
from openjarvis.connectors.pipeline import IngestionPipeline
from openjarvis.connectors.store import KnowledgeStore
from openjarvis.connectors.sync_engine import SyncEngine
cp = SyncEngine(
pipeline=IngestionPipeline(store=KnowledgeStore()),
).get_checkpoint(connector_id)
if cp and cp.get("items_synced") is not None:
baseline_items = int(cp["items_synced"])
except Exception: # noqa: BLE001
baseline_items = 0
_sync_state[connector_id] = {
"state": "syncing",
"error": None,
"baseline_items": baseline_items,
}
def _run_sync() -> None:
try:
from openjarvis.connectors.pipeline import IngestionPipeline
from openjarvis.connectors.store import KnowledgeStore
from openjarvis.connectors.sync_engine import SyncEngine
store = KnowledgeStore()
pipeline = IngestionPipeline(store=store)
engine = SyncEngine(pipeline=pipeline)
engine.sync(instance)
logger.info("Sync completed for %s", connector_id)
_sync_state[connector_id] = {
"state": "complete",
"error": None,
"baseline_items": baseline_items,
}
except Exception as exc:
error_msg = _translate_sync_error(str(exc))
logger.error("Sync failed for %s: %s", connector_id, error_msg)
_sync_state[connector_id] = {
"state": "error",
"error": error_msg,
"baseline_items": baseline_items,
}
t = threading.Thread(target=_run_sync, daemon=True)
t.start()
_sync_threads[connector_id] = t
return "started"
# ------------------------------------------------------------------
# Endpoints
# ------------------------------------------------------------------
@@ -248,39 +337,19 @@ def create_connectors_router():
except Exception as exc:
raise HTTPException(status_code=400, detail=str(exc))
# Auto-ingest after successful connection
# Auto-trigger a full backfill on a successful connect. Routed
# through the same _start_sync helper that POST /sync uses so the
# connection's progress is visible to GET /{id}/sync polling
# immediately — the user shouldn't have to click "Sync Now".
sync_status: Optional[str] = None
if instance.is_connected():
import threading
def _ingest() -> None:
try:
from openjarvis.connectors.pipeline import (
IngestionPipeline,
)
from openjarvis.connectors.store import KnowledgeStore
from openjarvis.connectors.sync_engine import SyncEngine
with KnowledgeStore() as store:
pipeline = IngestionPipeline(store)
engine = SyncEngine(pipeline)
engine.sync(instance)
logger.info(
"Auto-ingested %s after connect",
connector_id,
)
except Exception as exc:
logger.warning(
"Auto-ingest failed for %s: %s",
connector_id,
exc,
)
threading.Thread(target=_ingest, daemon=True).start()
sync_status = _start_sync(connector_id, instance)
return {
"connector_id": connector_id,
"connected": instance.is_connected(),
"status": "connected" if instance.is_connected() else "pending",
"sync_status": sync_status,
}
@router.post("/{connector_id}/disconnect")
@@ -440,15 +509,9 @@ def create_connectors_router():
)
)
# Track background sync state per connector
_sync_threads: Dict[str, Any] = {}
_sync_state: Dict[str, Dict[str, Any]] = {} # {connector_id: {state, error}}
@router.post("/{connector_id}/sync")
def trigger_sync(connector_id: str) -> Dict[str, Any]:
"""Trigger a sync in the background and return immediately."""
import threading
_ensure_connectors_registered()
if not ConnectorRegistry.contains(connector_id):
raise HTTPException(
@@ -461,51 +524,8 @@ def create_connectors_router():
status_code=400,
detail=f"Connector '{connector_id}' is not connected",
)
# If already syncing, don't start another
existing = _sync_threads.get(connector_id)
if existing and existing.is_alive():
return {
"connector_id": connector_id,
"status": "already_syncing",
}
# Mark as syncing immediately so the UI picks it up
_sync_state[connector_id] = {"state": "syncing", "error": None}
def _run_sync() -> None:
try:
from openjarvis.connectors.pipeline import IngestionPipeline
from openjarvis.connectors.store import KnowledgeStore
from openjarvis.connectors.sync_engine import SyncEngine
with KnowledgeStore() as store:
pipeline = IngestionPipeline(store=store)
engine = SyncEngine(pipeline=pipeline)
engine.sync(inst)
logger.info("Sync completed for %s", connector_id)
_sync_state[connector_id] = {"state": "complete", "error": None}
except Exception as exc:
error_msg = str(exc)
if "401" in error_msg or "Unauthorized" in error_msg:
error_msg = "Authentication failed — credentials may have expired."
elif "403" in error_msg or "Forbidden" in error_msg:
error_msg = "Permission denied — check API scopes."
elif "429" in error_msg or "Too Many Requests" in error_msg:
error_msg = "Rate limited — wait a minute and try again."
elif "timeout" in error_msg.lower():
error_msg = "Connection timed out."
logger.error("Sync failed for %s: %s", connector_id, error_msg)
_sync_state[connector_id] = {"state": "error", "error": error_msg}
t = threading.Thread(target=_run_sync, daemon=True)
t.start()
_sync_threads[connector_id] = t
return {
"connector_id": connector_id,
"status": "started",
}
status = _start_sync(connector_id, inst)
return {"connector_id": connector_id, "status": status}
@router.get("/{connector_id}/sync")
async def sync_status(connector_id: str):
@@ -527,6 +547,68 @@ def create_connectors_router():
bg_thread = _sync_threads.get(connector_id)
is_bg_running = bg_thread is not None and bg_thread.is_alive()
# Fall back to the SyncEngine's persistent checkpoint for items_synced
# and last_sync. The connector's own sync_status() only reflects state
# accumulated since the current connector instance was created, so a
# fresh process / fresh `_get_or_create` flips back to zeros even when
# tens of thousands of items have already been ingested historically.
# The checkpoint table is the source of truth.
checkpoint: Optional[Dict[str, Any]] = None
oldest_item_date: Optional[str] = None
try:
from openjarvis.connectors.pipeline import IngestionPipeline
from openjarvis.connectors.store import KnowledgeStore
from openjarvis.connectors.sync_engine import SyncEngine
store = KnowledgeStore()
checkpoint = SyncEngine(
pipeline=IngestionPipeline(store=store),
).get_checkpoint(connector_id)
# Map connector_id → source field as written by the connector.
# Most match 1:1, but the IMAP/OAuth Gmail connectors both write
# source='gmail' so the unified card pulls a single timeline.
_STORE_SOURCE = {"gmail_imap": "gmail"}
store_source = _STORE_SOURCE.get(connector_id, connector_id)
# Oldest indexed item for this connector so the UI can show how
# far back the corpus reaches ("synced back to 2024-03-12").
# ISO 8601 strings sort correctly under MIN(); skip rows with no
# timestamp and any tombstoned chunks so the display reflects
# what's actually retrievable.
row = store._conn.execute(
"SELECT MIN(timestamp) FROM knowledge_chunks "
"WHERE source = ? AND timestamp != '' AND deleted_at IS NULL",
(store_source,),
).fetchone()
if row is not None and row[0]:
oldest_item_date = str(row[0])
except Exception: # noqa: BLE001
checkpoint = None
# Prefer the checkpoint's items_synced — it is the running cumulative
# total across every sync this connector has ever run. The connector
# instance's own counter only tracks what *this* in-memory run has
# processed, so after a fresh sync of 30 new emails the connector
# reports 30 while the checkpoint reflects the correct 8,626 total.
if checkpoint and checkpoint.get("items_synced") is not None:
items_synced = int(checkpoint["items_synced"])
else:
items_synced = status.items_synced
# items_total reflects the connector's currently-known target for
# this sync run (e.g. number of IMAP message IDs matched). Leave it
# at 0 when unknown — the UI uses items_synced >= items_total > 0 as
# the "complete inbox" signal, which would spuriously fire on every
# page load if we masked the zero with items_synced as a fallback.
items_total = status.items_total
last_sync_str: Optional[str]
if status.last_sync:
last_sync_str = status.last_sync.isoformat()
elif checkpoint and checkpoint.get("last_sync"):
last_sync_str = checkpoint["last_sync"]
else:
last_sync_str = None
# Determine effective state
if is_bg_running:
effective_state = "syncing"
@@ -538,14 +620,27 @@ def create_connectors_router():
effective_state = status.state
# Use the bg error if the connector doesn't have one
effective_error = status.error or bg.get("error")
effective_error = (
status.error or bg.get("error") or (checkpoint or {}).get("error")
)
# Items processed during the *current* (or just-finished) run = total
# minus the baseline captured when this sync was kicked off. This
# lets the UI surface "30 new emails (8,623 total indexed)" without
# the client having to track its own baseline.
baseline_items = bg.get("baseline_items")
new_items_synced: Optional[int] = None
if isinstance(baseline_items, int):
new_items_synced = max(0, items_synced - baseline_items)
return {
"connector_id": connector_id,
"state": effective_state,
"items_synced": status.items_synced,
"items_total": status.items_total,
"last_sync": (status.last_sync.isoformat() if status.last_sync else None),
"items_synced": items_synced,
"items_total": items_total,
"new_items_synced": new_items_synced,
"last_sync": last_sync_str,
"oldest_item_date": oldest_item_date,
"error": effective_error,
}
+478
View File
@@ -0,0 +1,478 @@
"""HTTP route: ``POST /api/research`` — agentic research over the knowledge store.
Drives :class:`openjarvis.agents.research_loop.ResearchAgent` and streams a
custom SSE event schema back to the client:
* ``search_call`` about to invoke ``HybridSearch.search`` (with arguments)
* ``search_result`` search returned (num_hits, top_titles)
* ``synthesis`` final answer, emitted in word-window chunks for an
incremental UX (the agent itself returns the full string in one shot;
chunking happens in the router so we don't need to rewire the loop)
* ``done`` sentinel marking the end of the stream
Clarify is **disabled for the web session** the agent's clarify_handler is
overridden to return a fixed "no clarification available" string so the
loop never blocks waiting for terminal stdin. Bringing real clarify back
to the browser will require a two-step session protocol; that's a future
endpoint, not this one.
"""
from __future__ import annotations
import asyncio
import json
import logging
import re
import threading
import time
from typing import Any, AsyncGenerator, Callable, Dict, Optional
from fastapi import APIRouter
from fastapi.responses import StreamingResponse
from pydantic import BaseModel, Field
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.core.config import DEFAULT_CONFIG_DIR
from openjarvis.core.types import TelemetryRecord
from openjarvis.engine.ollama import OllamaEngine
from openjarvis.telemetry.store import TelemetryStore
logger = logging.getLogger(__name__)
router = APIRouter(prefix="/api", tags=["research"])
_WEB_CLARIFY_RESPONSE = "no clarification available in web session"
# Sentinel placed on the queue when the agent thread terminates.
_DONE = object()
def _record_research_telemetry(
*,
model: str,
usage: Dict[str, int],
latency_seconds: float,
energy_joules: float = 0.0,
mean_power_watts: float = 0.0,
peak_power_watts: float = 0.0,
energy_method: str = "",
) -> None:
"""Persist a research run into the telemetry DB so /v1/savings includes it.
GPU energy/power are sampled by ``_LiveGPUSampler`` during ``agent.run()``
and passed through here. With these fields populated, the same
``/v1/telemetry/energy`` aggregation that powers the chat System panel
rolls research into the same numbers Power (W) and Energy (kJ) stop
reading 0.0 after a research-only session.
Failures are swallowed telemetry persistence is best-effort and must
never break the user-visible SSE stream.
"""
if not usage:
return
db_path = DEFAULT_CONFIG_DIR / "telemetry.db"
try:
store = TelemetryStore(db_path)
except Exception as exc: # noqa: BLE001
logger.debug("research telemetry: cannot open %s: %s", db_path, exc)
return
try:
rec = TelemetryRecord(
timestamp=time.time(),
model_id=model,
engine="ollama",
agent="research",
prompt_tokens=int(usage.get("prompt_tokens", 0)),
prompt_tokens_evaluated=int(usage.get("prompt_tokens", 0)),
completion_tokens=int(usage.get("completion_tokens", 0)),
total_tokens=int(usage.get("total_tokens", 0)),
latency_seconds=latency_seconds,
energy_joules=energy_joules,
gpu_energy_joules=energy_joules,
power_watts=mean_power_watts,
energy_method=energy_method,
energy_vendor="nvidia" if energy_method else "",
is_streaming=True,
metadata={"peak_power_w": peak_power_watts} if peak_power_watts else {},
)
store.record(rec)
except Exception as exc: # noqa: BLE001
logger.debug("research telemetry: failed to record: %s", exc)
finally:
try:
store.close()
except Exception: # noqa: BLE001
pass
# ---------------------------------------------------------------------------
# Live GPU sampler — streams power/energy events during agent.run()
# ---------------------------------------------------------------------------
class _LiveGPUSampler:
"""Background pynvml poller — emits live power and accumulated energy.
Runs a small daemon thread for the lifetime of a research query. Every
``interval_s`` seconds it samples instantaneous GPU power across all
visible devices, rectangle-integrates it into a running energy total, and
invokes ``on_sample`` so the SSE consumer can forward ``system_metrics``
frames to the browser in real time.
On systems without pynvml (or with no visible GPU), ``available`` is
False and ``start()`` / ``stop()`` are no-ops research still runs, the
System panel just won't see live metrics for that session.
"""
def __init__(
self,
on_sample: Callable[[float, float, float], None],
interval_s: float = 0.75,
) -> None:
self._on_sample = on_sample
self._interval_s = interval_s
self._handles: list = []
self._pynvml = None
self._available = False
self._stop = threading.Event()
self._thread: Optional[threading.Thread] = None
self._energy_j = 0.0
self._peak_w = 0.0
self._power_sum = 0.0
self._sample_count = 0
self._t_start = 0.0
self._t_last = 0.0
try:
import pynvml # type: ignore
pynvml.nvmlInit()
count = pynvml.nvmlDeviceGetCount()
self._handles = [
pynvml.nvmlDeviceGetHandleByIndex(i) for i in range(count)
]
self._pynvml = pynvml
self._available = bool(self._handles)
if not self._available:
pynvml.nvmlShutdown()
except Exception as exc: # noqa: BLE001
logger.debug("research: pynvml unavailable, no live GPU metrics: %s", exc)
self._available = False
@property
def available(self) -> bool:
return self._available
def _poll_power_w(self) -> float:
total = 0.0
for h in self._handles:
try:
total += self._pynvml.nvmlDeviceGetPowerUsage(h) / 1000.0
except Exception: # noqa: BLE001
pass
return total
def _loop(self) -> None:
while not self._stop.is_set():
now = time.monotonic()
p = self._poll_power_w()
dt = now - self._t_last
# Rectangle integration is fine at sub-second cadence; the chat
# path uses the same approach in NvidiaEnergyMonitor's fallback.
self._energy_j += p * dt
self._t_last = now
self._sample_count += 1
self._power_sum += p
if p > self._peak_w:
self._peak_w = p
try:
self._on_sample(p, self._energy_j, now - self._t_start)
except Exception as exc: # noqa: BLE001
logger.debug("research: on_sample callback raised: %s", exc)
self._stop.wait(self._interval_s)
def start(self) -> None:
if not self._available:
return
self._t_start = time.monotonic()
self._t_last = self._t_start
self._thread = threading.Thread(target=self._loop, daemon=True)
self._thread.start()
def stop(self) -> Dict[str, float]:
if not self._available:
return {
"energy_j": 0.0,
"mean_power_w": 0.0,
"peak_power_w": 0.0,
"duration_s": 0.0,
}
self._stop.set()
if self._thread is not None:
self._thread.join(timeout=2.0)
duration = time.monotonic() - self._t_start
mean = self._power_sum / self._sample_count if self._sample_count else 0.0
try:
self._pynvml.nvmlShutdown()
except Exception: # noqa: BLE001
pass
return {
"energy_j": self._energy_j,
"mean_power_w": mean,
"peak_power_w": self._peak_w,
"duration_s": duration,
}
# ---------------------------------------------------------------------------
# Request model
# ---------------------------------------------------------------------------
class ResearchRequest(BaseModel):
query: str = Field(..., description="Natural-language question to research.")
# Deep Research has its own model requirements (function-calling support,
# sufficient reasoning capability) that the chat-model selector should not
# override. We accept the field for forward-compat with older clients but
# ignore it — the planner always runs on DEFAULT_PLANNER_MODEL.
model: Optional[str] = Field(
default=None, description="Ignored; retained for client compatibility."
)
# ---------------------------------------------------------------------------
# SSE helpers
# ---------------------------------------------------------------------------
def _sse(event: Dict[str, Any]) -> str:
"""Serialize one event dict to an SSE ``data: ...\\n\\n`` frame."""
return f"data: {json.dumps(event, ensure_ascii=False)}\n\n"
def _chunk_synthesis(text: str, window_chars: int = 40) -> list[str]:
"""Slice synthesis text into client-streaming-friendly chunks.
We break on word boundaries so partial deltas always render cleanly in
the browser. Each chunk is roughly ``window_chars`` characters long.
"""
if not text:
return []
tokens = re.findall(r"\S+\s*", text)
chunks: list[str] = []
buf = ""
for tok in tokens:
if len(buf) + len(tok) > window_chars and buf:
chunks.append(buf)
buf = tok
else:
buf += tok
if buf:
chunks.append(buf)
return chunks
# ---------------------------------------------------------------------------
# Stream generator
# ---------------------------------------------------------------------------
async def _stream_research(query: str, model: str) -> AsyncGenerator[str, None]:
"""Drive ResearchAgent on a worker thread; yield SSE frames as they land.
Three error envelopes setup, worker, consumer all funnel into the
same two-frame contract: ``{"type": "error", ...}`` followed by
``{"type": "done", "usage": {...}}``. The client can rely on always
seeing a ``done`` frame, even when the agent never started.
"""
# Phase 1: setup. Failures here (Ollama daemon down, DB locked, etc.)
# yield error + done and return — nothing has been emitted yet so the
# client gets a clean two-frame stream instead of a dangling connection.
try:
loop = asyncio.get_event_loop()
queue: asyncio.Queue = asyncio.Queue()
def on_event(event: Dict[str, Any]) -> None:
# Called from the agent's worker thread; bounce onto the event loop.
loop.call_soon_threadsafe(queue.put_nowait, event)
# Each request gets its own thin set of connectors. Constructing them
# is cheap (SQLite open + HTTP keepalive) and avoids state leaks
# between concurrent requests.
store = KnowledgeStore()
embedder = OllamaEmbedder()
if not embedder.is_available():
logger.warning(
"research: Ollama embedder unavailable; BM25-only retrieval."
)
embedder = None
engine = OllamaEngine()
agent = ResearchAgent(
engine=engine,
search=HybridSearch(store, embedder),
model=model,
clarify_handler=lambda question: _WEB_CLARIFY_RESPONSE,
on_event=on_event,
)
except Exception as exc: # noqa: BLE001
logger.exception("research: setup failed before agent could run: %s", exc)
yield _sse(
{
"type": "error",
"message": f"Research failed: {type(exc).__name__}: {exc}",
}
)
yield _sse({"type": "done", "usage": {}})
return
def _emit_live_sample(power_w: float, energy_j: float, duration_s: float) -> None:
# Called from the sampler's worker thread — bounce onto the asyncio
# loop so the SSE consumer sees the same queue ordering as agent
# events. The frontend mirrors these into the System panel so Power
# (W) and Energy (kJ) update live during the run.
loop.call_soon_threadsafe(
queue.put_nowait,
{
"type": "system_metrics",
"power_w": round(power_w, 2),
"energy_j": round(energy_j, 2),
"duration_s": round(duration_s, 2),
},
)
sampler = _LiveGPUSampler(on_sample=_emit_live_sample)
def _run() -> None:
t0 = time.time()
sampler.start()
try:
result = agent.run(query)
usage_dict = dict(result.usage)
totals = sampler.stop()
# Persist token usage *and* GPU energy/power so /v1/telemetry/energy
# rolls research into the same Power/Energy numbers as chat —
# this is what the launch-video System panel reads.
_record_research_telemetry(
model=model,
usage=usage_dict,
latency_seconds=time.time() - t0,
energy_joules=totals["energy_j"],
mean_power_watts=totals["mean_power_w"],
peak_power_watts=totals["peak_power_w"],
energy_method="polling" if sampler.available else "",
)
# Forward the aggregated token usage so the consumer can attach it
# to the terminal `done` frame. Internal event type — never sent to
# the client directly.
loop.call_soon_threadsafe(
queue.put_nowait,
{"type": "_usage", "usage": usage_dict},
)
except Exception as exc: # noqa: BLE001
logger.exception("research agent crashed: %s", exc)
# Stop the sampler on failure too so we don't leak the polling thread
# past the request lifetime.
try:
sampler.stop()
except Exception: # noqa: BLE001
pass
loop.call_soon_threadsafe(
queue.put_nowait,
{"type": "error", "message": f"{type(exc).__name__}: {exc}"},
)
finally:
loop.call_soon_threadsafe(queue.put_nowait, _DONE)
task = asyncio.create_task(asyncio.to_thread(_run))
final_answer: Optional[str] = None
final_usage: Dict[str, int] = {}
try:
while True:
event = await queue.get()
if event is _DONE:
break
if not isinstance(event, dict):
continue
etype = event.get("type")
# Internal usage marker — capture for the done frame, don't forward.
if etype == "_usage":
final_usage = event.get("usage", {}) or {}
continue
# We translate the agent's `final_answer` event into a stream of
# `synthesis` chunks so the client sees the answer materialize
# incrementally rather than as a single blob.
if etype == "final_answer":
final_answer = event.get("text", "")
for piece in _chunk_synthesis(final_answer or ""):
yield _sse({"type": "synthesis", "text": piece})
continue
yield _sse(event)
# If the agent thread crashed before producing a final answer, the
# client still gets the error frame (emitted above) followed by done.
yield _sse({"type": "done", "usage": final_usage})
except Exception as exc: # noqa: BLE001
# Consumer loop crashed unexpectedly (e.g. JSON serialization fault,
# logic bug). Surface a clean error frame rather than letting the
# SSE connection die mid-stream.
logger.exception("research: stream consumer crashed: %s", exc)
yield _sse(
{
"type": "error",
"message": f"Research failed: {type(exc).__name__}: {exc}",
}
)
yield _sse({"type": "done", "usage": final_usage})
finally:
# The worker may still be cleaning up (rarely) — make sure we don't
# leak a dangling task. Swallow any straggler exception so a worker
# failure during teardown doesn't escape the generator after we've
# already emitted the terminal done frame.
if not task.done():
try:
await task
except Exception as exc: # noqa: BLE001
logger.debug("research: worker task ended with %s", exc)
# ---------------------------------------------------------------------------
# Route
# ---------------------------------------------------------------------------
@router.post("/research")
async def research(req: ResearchRequest) -> StreamingResponse:
"""Run a research query and stream the agent's trace + synthesis via SSE.
Response is ``text/event-stream`` with one JSON event per frame. See the
module docstring for the schema; a final ``{"type": "done"}`` always
terminates the stream so clients can detect end-of-response without
parsing the underlying ``[DONE]`` sentinel used by OpenAI-style routes.
"""
if req.model and req.model != DEFAULT_PLANNER_MODEL:
logger.info(
"research: ignoring client model=%r; using DEFAULT_PLANNER_MODEL=%r",
req.model,
DEFAULT_PLANNER_MODEL,
)
return StreamingResponse(
_stream_research(req.query, DEFAULT_PLANNER_MODEL),
media_type="text/event-stream",
headers={
"Cache-Control": "no-cache",
"Connection": "keep-alive",
"X-Accel-Buffering": "no",
},
)
__all__ = ["router", "ResearchRequest"]
+186
View File
@@ -0,0 +1,186 @@
"""Tests for the research_loop agent — focused on loop invariants.
The fixtures use a minimal mock engine so these tests are fast and don't
require a running Ollama daemon.
"""
from __future__ import annotations
import json
from typing import Any, Dict, List, Optional, Sequence
from unittest.mock import MagicMock
import pytest
from openjarvis.agents.research_loop import ResearchAgent
class _MockEngine:
"""Engine stub that lets a test script the per-call response.
``responses`` is a list of dicts in ``OllamaEngine.generate`` shape
(``content``, ``tool_calls``, ``usage``); each call to ``generate``
consumes the next entry. When the script runs out, the engine repeats
the final entry handy for "loop forever" mock behaviors.
"""
def __init__(self, responses: List[Dict[str, Any]]):
self._responses = list(responses)
self.calls: List[Dict[str, Any]] = []
def generate(
self,
messages: Sequence,
*,
model: str,
temperature: float = 0.0,
max_tokens: int = 1024,
num_ctx: int = 8192,
tools: Optional[List[Dict[str, Any]]] = None,
**kwargs: Any,
) -> Dict[str, Any]:
self.calls.append({"tools": tools, "messages": list(messages)})
if self._responses:
return (
self._responses.pop(0)
if len(self._responses) > 1
else self._responses[0]
)
return {"content": "", "tool_calls": [], "usage": {}}
def _search_call(call_id: str, query: str = "anything") -> Dict[str, Any]:
return {
"content": "",
"tool_calls": [
{
"id": call_id,
"name": "search",
"arguments": json.dumps({"query": query}),
}
],
"usage": {"prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15},
}
def _text_response(text: str) -> Dict[str, Any]:
return {
"content": text,
"tool_calls": [],
"usage": {"prompt_tokens": 20, "completion_tokens": 10, "total_tokens": 30},
}
@pytest.fixture()
def stub_search() -> MagicMock:
"""A HybridSearch stand-in whose .search() returns an empty hit list."""
s = MagicMock()
s.search.return_value = []
return s
def test_forced_synthesis_when_budget_exhausts(stub_search: MagicMock) -> None:
"""Loop exits cleanly with a synthesis even when the model keeps tool-calling.
With max_iterations=1, the budget is exhausted after a single search.
The next engine call receives ``tools=None`` and is expected to produce
a text answer; the new forced-synthesis fallback then makes one more
no-tools call to ensure we always return something usable.
"""
# Force the loop to fall through to the new forced-synthesis path: every
# in-loop iteration returns a tool call (so we never take an early text
# return), and the final post-loop call returns text.
engine = _MockEngine(
responses=[
_search_call("call-1"),
_search_call("call-2"),
_text_response("Here is what I found: nothing usable."),
]
)
agent = ResearchAgent(engine, stub_search, model="mock", max_iterations=1)
result = agent.run("test query")
assert "Here is what I found" in result.answer
assert "loop exhausted" not in result.answer.lower()
# The final call was made with tools=None (the forced-synthesis path).
assert engine.calls[-1]["tools"] is None
assert all(t.tool_name == "search" for t in result.tool_calls)
def test_forced_synthesis_returns_sentinel_when_model_stays_silent(
stub_search: MagicMock,
) -> None:
"""If even the forced final call returns no text, surface a clear sentinel.
This is the gracefully-degraded case: the loop did its best but the model
refused to emit text. The contract is still that ``answer`` is a non-empty
string the caller can show to the user.
"""
engine = _MockEngine(
responses=[
_search_call("call-1"),
_search_call("call-2"),
{"content": "", "tool_calls": [], "usage": {}}, # silent forced call
]
)
agent = ResearchAgent(engine, stub_search, model="mock", max_iterations=1)
result = agent.run("test query")
assert result.answer # non-empty
assert "no synthesis available" in result.answer.lower()
def test_first_turn_text_response_returns_directly(stub_search: MagicMock) -> None:
"""When the model produces text on the very first turn, return it as-is."""
engine = _MockEngine(responses=[_text_response("Quick answer.")])
agent = ResearchAgent(engine, stub_search, model="mock", max_iterations=5)
result = agent.run("hello")
assert result.answer == "Quick answer."
assert result.tool_calls == []
# Only one engine call needed.
assert len(engine.calls) == 1
def test_clarify_before_any_search_is_rejected(stub_search: MagicMock) -> None:
"""clarify() invoked before a search call must return a runtime error result.
The system prompt asks the planner to search first, but the loop also
enforces this at runtime so a non-compliant planner can't surprise the
user with a pre-emptive clarification.
"""
clarify_calls: List[str] = []
def fake_clarify(q: str) -> str:
clarify_calls.append(q)
return "user response"
engine = _MockEngine(
responses=[
{
"content": "",
"tool_calls": [
{
"id": "c1",
"name": "clarify",
"arguments": json.dumps({"question": "who?"}),
}
],
"usage": {},
},
_text_response("Final."),
]
)
agent = ResearchAgent(
engine, stub_search, model="mock", max_iterations=5,
clarify_handler=fake_clarify,
)
result = agent.run("vague query")
# Handler must not have been called because the dispatch returned an error.
assert clarify_calls == []
# No clarify invocation recorded.
assert all(t.tool_name != "clarify" for t in result.tool_calls)
assert result.answer == "Final."
+494 -2
View File
@@ -5,6 +5,7 @@ All Gmail API calls are mocked; no network access is required.
from __future__ import annotations
import base64
import json
from pathlib import Path
from typing import List
@@ -225,7 +226,8 @@ def test_sync_passes_since_as_query(
# Verify the epoch value is correct
expected_epoch = int(since_dt.timestamp())
assert f"after:{expected_epoch}" in call_kwargs["query"]
assert "category:primary" in call_kwargs["query"]
# No category filter — sent mail and all Gmail tabs should be reachable.
assert "category:" not in call_kwargs["query"]
# ---------------------------------------------------------------------------
@@ -251,7 +253,8 @@ def test_sync_without_since_passes_empty_query(
mock_list.assert_called_once()
_, call_kwargs = mock_list.call_args
assert call_kwargs.get("query", "") == "category:primary"
# No since= and no hardcoded category filter → empty query.
assert call_kwargs.get("query", "") == ""
# ---------------------------------------------------------------------------
@@ -269,3 +272,492 @@ def test_registry() -> None:
assert ConnectorRegistry.contains("gmail")
cls = ConnectorRegistry.get("gmail")
assert cls.connector_id == "gmail"
# ---------------------------------------------------------------------------
# v1 schema tests
# ---------------------------------------------------------------------------
# base64url("Migration body") == "TWlncmF0aW9uIGJvZHk="
_MSG_RICH = {
"id": "msg-rich-1",
"threadId": "thread-rich-1",
"labelIds": ["SENT", "CATEGORY_PERSONAL"],
"snippet": "Quick note on the migration",
"historyId": "1234567",
"sizeEstimate": 4096,
"payload": {
"mimeType": "text/plain",
"headers": [
{"name": "From", "value": "Alice Bose <alice@example.com>"},
{
"name": "To",
"value": "Bob <bob@example.com>, charlie@example.com",
},
{"name": "Cc", "value": "dana@example.com"},
{"name": "Subject", "value": "Migration plan"},
{"name": "Date", "value": "Mon, 01 Jan 2024 10:00:00 +0000"},
{"name": "Message-ID", "value": "<abc123@mail.example.com>"},
],
"body": {"data": "TWlncmF0aW9uIGJvZHk="},
},
}
@patch("openjarvis.connectors.gmail._gmail_api_list_messages")
@patch("openjarvis.connectors.gmail._gmail_api_get_message")
def test_sync_emits_v1_fields(
mock_get,
mock_list,
connector,
tmp_path: Path,
) -> None:
"""sync() populates source_id, participants_raw, normalized participants,
channel from labels, and Gmail-specific fields in metadata."""
creds_path = Path(connector._credentials_path)
creds_path.write_text(json.dumps({"token": "fake-access-token"}), encoding="utf-8")
mock_list.return_value = {"messages": [{"id": "msg-rich-1"}]}
mock_get.return_value = _MSG_RICH
docs: List[Document] = list(connector.sync())
assert len(docs) == 1
doc = docs[0]
# source_id is the raw Gmail msg.id (no "gmail:" prefix); doc_id stays
# composite for back-compat with legacy callers.
assert doc.source_id == "msg-rich-1"
assert doc.doc_id == "gmail:msg-rich-1"
# participants_raw preserves From / To / Cc as Gmail returned them.
assert doc.participants_raw == [
"Alice Bose <alice@example.com>",
"Bob <bob@example.com>, charlie@example.com",
"dana@example.com",
]
# participants extracts and lowercases all addresses, multi-recipient-aware.
assert doc.participants == [
"alice@example.com",
"bob@example.com",
"charlie@example.com",
"dana@example.com",
]
# channel comes from the highest-priority system label.
assert doc.channel == "SENT"
# Gmail-specific richness lives in source_metadata.
assert doc.metadata["rfc_message_id"] == "<abc123@mail.example.com>"
assert doc.metadata["snippet"] == "Quick note on the migration"
assert doc.metadata["history_id"] == "1234567"
assert doc.metadata["size_estimate"] == 4096
assert "SENT" in doc.metadata["labels"]
assert "CATEGORY_PERSONAL" in doc.metadata["labels"]
@patch("openjarvis.connectors.gmail._gmail_api_list_messages")
@patch("openjarvis.connectors.gmail._gmail_api_get_message")
def test_sync_channel_falls_back_to_inbox(
mock_get,
mock_list,
connector,
tmp_path: Path,
) -> None:
"""When only INBOX is present among system labels, channel == 'INBOX'."""
creds_path = Path(connector._credentials_path)
creds_path.write_text(json.dumps({"token": "fake-access-token"}), encoding="utf-8")
mock_list.return_value = _LIST_RESPONSE
mock_get.side_effect = lambda token, msg_id: _MSG1 if msg_id == "msg1" else _MSG2
docs = list(connector.sync())
assert all(d.channel == "INBOX" for d in docs)
@patch("openjarvis.connectors.gmail._gmail_api_list_messages")
@patch("openjarvis.connectors.gmail._gmail_api_get_message")
def test_sync_channel_none_when_no_system_label(
mock_get,
mock_list,
connector,
tmp_path: Path,
) -> None:
"""A message with only user-defined labels has channel=None."""
creds_path = Path(connector._credentials_path)
creds_path.write_text(json.dumps({"token": "fake-access-token"}), encoding="utf-8")
msg_userlabel = {
**_MSG1,
"id": "msg-user",
"labelIds": ["Label_1234", "CATEGORY_FORUMS"],
}
mock_list.return_value = {"messages": [{"id": "msg-user"}]}
mock_get.return_value = msg_userlabel
docs = list(connector.sync())
assert len(docs) == 1
assert docs[0].channel is None
# ---------------------------------------------------------------------------
# HTML body stripping
# ---------------------------------------------------------------------------
def test_html_to_text_strips_basic_tags() -> None:
"""_html_to_text() removes tags but preserves visible text content."""
from openjarvis.connectors.gmail import _html_to_text # noqa: PLC0415
html = (
"<html><body><p>Hello <b>world</b>!</p>"
"<p>Second paragraph.</p></body></html>"
)
text = _html_to_text(html)
assert "Hello" in text
assert "world" in text
assert "Second paragraph" in text
assert "<" not in text and ">" not in text
def test_html_to_text_drops_script_and_style() -> None:
"""Content inside <script>/<style>/<head> is stripped, not rendered."""
from openjarvis.connectors.gmail import _html_to_text # noqa: PLC0415
html = """
<html>
<head><style>.foo { color: red; }</style><title>Ignore me</title></head>
<body>
<script>alert('xss')</script>
<p>Visible content here.</p>
</body>
</html>
"""
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 = "<p>Tom &amp; Jerry &mdash; 50&nbsp;cents</p>"
text = _html_to_text(html)
assert "Tom & Jerry" in text
assert "" in text # &mdash;
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 = "<div>line one</div><div>line two</div><div>line three</div>"
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"<html><body><p>Hello <b>world</b>!</p>"
b"<p>Second paragraph.</p></body></html>"
)
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><body><p>HTML version</p></body></html>"
).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"
+1 -1
View File
@@ -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)
+171
View File
@@ -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] == ""
+106
View File
@@ -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 <alice@example.com>"],
)
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 <alice@example.com>"]
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
Generated
+2
View File
@@ -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" },