mirror of
https://github.com/open-jarvis/OpenJarvis.git
synced 2026-07-27 21:05:34 +00:00
feat: Deep Research — personal deep research over Gmail with hybrid retrieval and agentic synthesis (#354)
This commit is contained in:
@@ -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
|
||||
|
||||
@@ -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/
|
||||
|
||||
@@ -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>
|
||||
|
||||
@@ -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={{
|
||||
|
||||
@@ -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>
|
||||
);
|
||||
}
|
||||
@@ -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>
|
||||
|
||||
@@ -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})` : ''}` });
|
||||
}
|
||||
|
||||
@@ -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)' }}>
|
||||
|
||||
@@ -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
@@ -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> {
|
||||
|
||||
@@ -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
@@ -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();
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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 }),
|
||||
|
||||
@@ -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'} →
|
||||
</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)' }}>
|
||||
|
||||
@@ -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' },
|
||||
|
||||
@@ -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 };
|
||||
|
||||
@@ -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',
|
||||
},
|
||||
},
|
||||
});
|
||||
|
||||
@@ -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 = [
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -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()
|
||||
|
||||
@@ -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",
|
||||
]
|
||||
@@ -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
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -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
|
||||
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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)))
|
||||
|
||||
|
||||
@@ -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",
|
||||
]
|
||||
@@ -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,
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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]
|
||||
|
||||
@@ -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",
|
||||
]
|
||||
@@ -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", "")
|
||||
|
||||
@@ -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"]
|
||||
@@ -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
|
||||
|
||||
|
||||
@@ -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"""
|
||||
|
||||
@@ -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",
|
||||
|
||||
@@ -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)
|
||||
|
||||
|
||||
@@ -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,
|
||||
}
|
||||
|
||||
|
||||
@@ -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"]
|
||||
@@ -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."
|
||||
@@ -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 & Jerry — 50 cents</p>"
|
||||
text = _html_to_text(html)
|
||||
assert "Tom & Jerry" in text
|
||||
assert "—" in text # —
|
||||
assert "50" in text and "cents" in text
|
||||
|
||||
|
||||
def test_html_to_text_inserts_paragraph_breaks() -> None:
|
||||
"""Block-level tags produce newlines so the chunker sees structure."""
|
||||
from openjarvis.connectors.gmail import _html_to_text # noqa: PLC0415
|
||||
|
||||
html = "<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"
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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] == ""
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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" },
|
||||
|
||||
Reference in New Issue
Block a user