/ src / api / aipolaris.ts
aipolaris.ts
  1  // Copyright (c) 2026 VPL Solutions. All rights reserved.
  2  // Licensed under the MIT License. See LICENSE for details.
  3  //
  4  // aiPolaris API client — ADR-011
  5  //
  6  // Connects Studio to the aiPolaris Planner → Retriever → Synthesizer DAG.
  7  // Base URL: VITE_AIPOLARIS_URL (never hardcoded — ADR-011).
  8  //
  9  // TODO: pass Authorization header once ADR-011 JWT wiring is complete.
 10  //       aiPolaris AUTH_ENABLED=false for this integration phase.
 11  
 12  import { config } from '../config';
 13  
 14  // ── Types ─────────────────────────────────────────────────────────────────────
 15  
 16  export type DagStage = 'planner' | 'retriever' | 'synthesizer';
 17  export type DagStageStatus = 'idle' | 'running' | 'done' | 'error';
 18  
 19  export interface StepRecord {
 20    node_name: string;
 21    input_hash: string;
 22    tool_calls: string[];
 23    output_hash: string;
 24    latency_ms: number;
 25    timestamp: string;
 26  }
 27  
 28  export interface TraceContext {
 29    trace_id: string;
 30    steps: StepRecord[];
 31  }
 32  
 33  export interface PolarisChunk {
 34    title: string;
 35    content: string;
 36    source: string;
 37    last_modified: string;
 38    reranker_score: number;
 39  }
 40  
 41  export interface PolarisCitation {
 42    title: string;
 43    source: string;
 44    excerpt: string;
 45  }
 46  
 47  // SSE event shapes from aiPolaris /query endpoint
 48  export interface PolarisTokenEvent {
 49    type: 'token';
 50    content: string;
 51  }
 52  
 53  export interface PolarisDoneEvent {
 54    type: 'done';
 55    citations: PolarisCitation[];
 56    trace_id: string;
 57    latency_ms: number;
 58  }
 59  
 60  export interface PolarisErrorEvent {
 61    type: 'error';
 62    message: string;
 63  }
 64  
 65  export type PolarisStreamEvent =
 66    | PolarisTokenEvent
 67    | PolarisDoneEvent
 68    | PolarisErrorEvent;
 69  
 70  // Enriched DAG node state for UI rendering
 71  export interface DagNodeState {
 72    stage: DagStage;
 73    status: DagStageStatus;
 74    label: string;
 75    sublabel: string;                // e.g. "2 sub-tasks" | "5 chunks" | "3 citations"
 76    latency_ms: number | null;
 77    // Expanded drawer content
 78    subTasks?: string[];             // Planner output
 79    chunks?: PolarisChunk[];         // Retriever output
 80    citations?: PolarisCitation[];   // Synthesizer output
 81    errorHint?: string;              // Actionable recovery hint on failure
 82  }
 83  
 84  export interface PolarisQueryResult {
 85    answer: string;
 86    citations: PolarisCitation[];
 87    trace_id: string;
 88    session_id: string;
 89    latency_ms: number;
 90    dagNodes: DagNodeState[];
 91  }
 92  
 93  export interface AiPolarisError extends Error {
 94    stage?: DagStage;
 95    hint?: string;
 96  }
 97  
 98  // ── Stage progress inference from SSE ────────────────────────────────────────
 99  
100  // aiPolaris SSE events don't include per-stage notifications in the current
101  // implementation — we infer stage from token arrival timing.
102  // When StepRecord data is available in the done event, we backfill latencies.
103  
104  function initialDagNodes(): DagNodeState[] {
105    return [
106      { stage: 'planner',     status: 'idle', label: 'Planner',     sublabel: 'Decomposing question',  latency_ms: null },
107      { stage: 'retriever',   status: 'idle', label: 'Retriever',   sublabel: 'Searching knowledge base', latency_ms: null },
108      { stage: 'synthesizer', status: 'idle', label: 'Synthesizer', sublabel: 'Assembling answer',     latency_ms: null },
109    ];
110  }
111  
112  // ── Error hints per stage ─────────────────────────────────────────────────────
113  
114  const STAGE_ERROR_HINTS: Record<DagStage, string> = {
115    planner:     'The Planner could not decompose this question — try rephrasing it more specifically.',
116    retriever:   'The Retriever found no relevant documents — try broadening or rephrasing the question.',
117    synthesizer: 'The Synthesizer could not produce an answer from the retrieved content — the knowledge base may not cover this topic.',
118  };
119  
120  // ── API client ────────────────────────────────────────────────────────────────
121  
122  export type StageUpdateCallback = (nodes: DagNodeState[]) => void;
123  
124  /**
125   * Stream a query through the aiPolaris DAG.
126   *
127   * Calls POST /query on the aiPolaris API with SSE streaming.
128   * Invokes onStageUpdate as each stage transitions (idle → running → done).
129   * Resolves with the full PolarisQueryResult on completion.
130   * Rejects with AiPolarisError on network or stage failure.
131   *
132   * ADR-011: BASE URL from VITE_AIPOLARIS_URL — never hardcoded.
133   * ADR-011: No Authorization header — AUTH_ENABLED=false in this phase.
134   */
135  export async function streamPolarisQuery(
136    query: string,
137    sessionId: string | null,
138    onToken: (token: string) => void,
139    onStageUpdate: StageUpdateCallback,
140    signal?: AbortSignal,
141  ): Promise<PolarisQueryResult> {
142    const baseUrl = config.aipolarisBaseUrl;
143    const nodes = initialDagNodes();
144  
145    // Planner starts immediately on request
146    nodes[0] = { ...nodes[0], status: 'running', sublabel: 'Planning...' };
147    onStageUpdate([...nodes]);
148  
149    const body: Record<string, unknown> = { query };
150    if (sessionId) body.session_id = sessionId;
151  
152    let response: Response;
153    try {
154      response = await fetch(`${baseUrl}/query`, {
155        method: 'POST',
156        headers: { 'Content-Type': 'application/json' },
157        body: JSON.stringify(body),
158        signal,
159      });
160    } catch {
161      const apiErr = new Error('Failed to connect to aiPolaris') as AiPolarisError;
162      apiErr.stage = 'planner';
163      apiErr.hint = STAGE_ERROR_HINTS.planner;
164      nodes[0] = { ...nodes[0], status: 'error', errorHint: STAGE_ERROR_HINTS.planner };
165      onStageUpdate([...nodes]);
166      throw apiErr;
167    }
168  
169    if (!response.ok) {
170      const detail = await response.json().catch(() => ({}));
171      const apiErr = new Error(detail?.detail || `aiPolaris error: ${response.status}`) as AiPolarisError;
172      apiErr.stage = 'planner';
173      apiErr.hint = STAGE_ERROR_HINTS.planner;
174      nodes[0] = { ...nodes[0], status: 'error', errorHint: STAGE_ERROR_HINTS.planner };
175      onStageUpdate([...nodes]);
176      throw apiErr;
177    }
178  
179    // Response headers carry session + trace IDs
180    const responseSessionId = response.headers.get('X-Session-Id') ?? (sessionId || crypto.randomUUID());
181    const responseTraceId   = response.headers.get('X-Trace-Id') ?? '';
182  
183    // Planner done → Retriever running (first token arrival = planner + retriever complete)
184    nodes[0] = { ...nodes[0], status: 'done', sublabel: 'Done', latency_ms: null };
185    nodes[1] = { ...nodes[1], status: 'running', sublabel: 'Retrieving...' };
186    onStageUpdate([...nodes]);
187  
188    const reader = response.body?.getReader();
189    if (!reader) {
190      throw new Error('No response body from aiPolaris');
191    }
192  
193    const decoder = new TextDecoder();
194    let buffer = '';
195    const answerTokens: string[] = [];
196    let citations: PolarisCitation[] = [];
197    let latency_ms = 0;
198    let firstToken = true;
199  
200    try {
201      while (true) {
202        const { done, value } = await reader.read();
203        if (done) break;
204  
205        buffer += decoder.decode(value, { stream: true });
206        const lines = buffer.split('\n');
207        buffer = lines.pop() ?? '';
208  
209        for (const line of lines) {
210          if (!line.startsWith('data: ')) continue;
211          const raw = line.slice(6).trim();
212          if (!raw) continue;
213  
214          let event: PolarisStreamEvent;
215          try {
216            event = JSON.parse(raw) as PolarisStreamEvent;
217          } catch {
218            continue;
219          }
220  
221          if (event.type === 'token') {
222            // First token: retriever done → synthesizer running
223            if (firstToken) {
224              firstToken = false;
225              nodes[1] = { ...nodes[1], status: 'done', sublabel: 'Done', latency_ms: null };
226              nodes[2] = { ...nodes[2], status: 'running', sublabel: 'Synthesizing...' };
227              onStageUpdate([...nodes]);
228            }
229            answerTokens.push(event.content);
230            onToken(event.content);
231  
232          } else if (event.type === 'done') {
233            citations = event.citations;
234            latency_ms = event.latency_ms;
235            nodes[2] = {
236              ...nodes[2],
237              status: 'done',
238              sublabel: `${citations.length} citation${citations.length !== 1 ? 's' : ''}`,
239              latency_ms: event.latency_ms,
240              citations,
241            };
242            onStageUpdate([...nodes]);
243  
244          } else if (event.type === 'error') {
245            // Determine which stage failed based on answer state
246            const failedStage: DagStage = firstToken ? 'retriever' : 'synthesizer';
247            const stageIdx = failedStage === 'retriever' ? 1 : 2;
248            nodes[stageIdx] = {
249              ...nodes[stageIdx],
250              status: 'error',
251              errorHint: STAGE_ERROR_HINTS[failedStage],
252            };
253            onStageUpdate([...nodes]);
254            const apiErr = new Error(event.message) as AiPolarisError;
255            apiErr.stage = failedStage;
256            apiErr.hint = STAGE_ERROR_HINTS[failedStage];
257            throw apiErr;
258          }
259        }
260      }
261    } finally {
262      reader.releaseLock();
263    }
264  
265    return {
266      answer: answerTokens.join(''),
267      citations,
268      trace_id: responseTraceId,
269      session_id: responseSessionId,
270      latency_ms,
271      dagNodes: nodes,
272    };
273  }
274  
275  /**
276   * Complexity heuristic for auto-routing hint.
277   * Returns true if the query looks like it would benefit from Agent Query.
278   * Used to show a suggestion banner — never silently switches modes.
279   */
280  export function isComplexQuery(query: string): boolean {
281    const q = query.toLowerCase().trim();
282    const multiStepPatterns = [
283      /\band\b.*\band\b/,           // "X and Y and Z"
284      /\bhow does .+ compare/,      // comparison
285      /\bwhat .+ and (who|when|where|how)\b/, // compound
286      /\bsteps?\b.*\band\b/,        // steps + conjunction
287      /\bdifference between\b/,
288      /\bpros and cons\b/,
289      /\bwalk me through\b/,
290      /\bexplain .+ and .+/,
291    ];
292    return multiStepPatterns.some(p => p.test(q));
293  }