@od-oneapp/ai-platform 0.1.0

package/src/grounding/rag/tool.ts

@@ -0,0 +1,281 @@
+/**
+ * @fileoverview Research tool for agent integration
+ *
+ * Wraps the three-stage research pipeline as an AI SDK tool
+ * that can be used by any agent (canvas, chat, etc.)
+ *
+ * Features:
+ * - Configurable models per stage via PipelineModelConfig
+ * - Tier-based research (fast/balanced/deep)
+ * - Full trace with cost/latency
+ * - Compliance signal integration
+ * - Artifact storage for research results
+ *
+ * @module @od-oneapp/ai-platform/research/tool
+ */
+
+import { z } from 'zod';
+
+import { tool } from '../../shared';
+
+
+import { reasonWithSonnet, researchWithPerplexity, verifyWithGPT5 } from './pipeline';
+import {
+  endRun,
+  endStep,
+  persistRunIndex,
+  persistRunMeta,
+  persistStep,
+  startRun,
+  startStep,
+} from './trace';
+
+import type { PipelineModelConfig, ResearchTier } from './types';
+import type { ComplianceSignal } from '../../governance/compliance/schemas';
+
+/**
+ * Schema for research tool input.
+ */
+export const researchInputSchema = z.object({
+  /** Research question or query. */
+  question: z.string().min(1).describe('Research question or query'),
+  /** Research depth: fast (400 tokens), balanced (800 tokens), deep (1600 tokens). */
+  tier: z
+    .enum(['fast', 'balanced', 'deep'])
+    .optional()
+    .default('balanced')
+    .describe('Research depth: fast (400 tokens), balanced (800 tokens), deep (1600 tokens)'),
+});
+
+/** Research tool input type. */
+export type ResearchInput = z.infer<typeof researchInputSchema>;
+
+/**
+ * Creates a research tool for AI SDK agents.
+ * Wraps the multi-model research pipeline with trace storage and compliance signals.
+ *
+ * @param options - Optional tool configuration:
+ * - `onSignal` - Callback for compliance signal emission
+ * - `onTrace` - Callback for trace step completion (receives runId, step kind, latency, cost)
+ * @returns AI SDK tool ready for agent integration.
+ * @remarks Tool executes full three-stage research pipeline (context expansion → reasoning → verification).
+ * Persists step artifacts and run metadata for traceability.
+ * Emits compliance signals at start and completion.
+ * Returns structured result with answer, confidence, citations, flags, and trace.
+ * On error, returns failure result with error message and emits guardrail.blocked signal.
+ * @example
+ * ```ts
+ * const researchTool = createResearchTool({
+ *   onSignal: (signal) => persistSignal(signal),
+ *   onTrace: (runId, step) => console.log(`Step ${step.kind} completed`)
+ * });
+ *
+ * const agent = createAgent({
+ *   tools: { research: researchTool }
+ * });
+ * ```
+ */
+export function createResearchTool(options?: {
+  onSignal?: (signal: ComplianceSignal) => void;
+  onTrace?: (runId: string, step: { kind: string; latencyMs: number; costUSD: number }) => void;
+  /** Optional model configuration overrides for the three-stage pipeline. */
+  modelConfig?: PipelineModelConfig;
+}) {
+  return tool({
+    description: `Perform multi-model research using a three-stage pipeline: context expansion → structured reasoning → verification.
+
+This tool provides grounded, verified research with:
+- Context expansion (notes + citations)
+- Structured reasoning (claims, assumptions, gaps)
+- Verification (fact-checking, confidence scoring)
+
+Use this tool when you need factual, citation-backed answers to complex questions.
+
+Tiers:
+- fast: Quick research (400 tokens, ~2-3 citations)
+- balanced: Standard research (800 tokens, ~5 citations)
+- deep: Comprehensive research (1600 tokens, ~7 citations)`,
+    inputSchema: researchInputSchema,
+    execute: async (input: ResearchInput) => {
+      const { question, tier } = input;
+
+      try {
+        // Start run trace
+        const run = startRun(question, tier as ResearchTier, 'next-ai');
+
+        // Emit start signal
+        options?.onSignal?.({
+          id: crypto.randomUUID(),
+          type: 'tool.executed',
+          message: `Starting multi-model research (tier: ${tier})`,
+          runId: run.runId,
+          timestamp: new Date().toISOString(),
+          metadata: { question, tier },
+        });
+
+        const mc = options?.modelConfig;
+
+        // STEP 1: Context expansion
+        const expansionModel = mc?.contextExpansion?.model ?? 'llama-3.1-sonar-large-128k-online';
+        const expansionStep = startStep('context-expansion', expansionModel);
+        const {
+          notes,
+          cites,
+          usage: expansionUsage,
+        } = await researchWithPerplexity(question, tier as ResearchTier);
+        endStep(expansionStep, {
+          inputTokens: expansionUsage.inputTokens,
+          outputTokens: expansionUsage.outputTokens,
+        });
+
+        await persistStep(run, expansionStep, {
+          expansion_notes: { notes, cites },
+          expansion_index: {
+            sources: cites,
+            fetchedAt: new Date().toISOString(),
+            tier,
+          },
+        });
+        run.steps.push(expansionStep);
+
+        options?.onTrace?.(run.runId, {
+          kind: 'context-expansion',
+          latencyMs: expansionStep.latencyMs,
+          costUSD: expansionStep.costUSD ?? 0,
+        });
+
+        // STEP 2: Structured reasoning
+        const reasoningModel = mc?.reasoning?.model ?? 'claude-3-5-sonnet-20241022';
+        const reasoningStep = startStep('reasoning', reasoningModel);
+        const reasoning = await reasonWithSonnet({ question, notes, cites });
+        endStep(reasoningStep, {
+          inputTokens: reasoning.usage.inputTokens,
+          outputTokens: reasoning.usage.outputTokens,
+        });
+
+        await persistStep(run, reasoningStep, {
+          reasoning_output: {
+            claims: reasoning.claims,
+            assumptions: reasoning.assumptions,
+            gaps: reasoning.gaps,
+            answerDraft: reasoning.answerDraft,
+            meta: {
+              model: reasoningStep.model,
+              timestamp: new Date().toISOString(),
+            },
+          },
+        });
+        run.steps.push(reasoningStep);
+
+        options?.onTrace?.(run.runId, {
+          kind: 'reasoning',
+          latencyMs: reasoningStep.latencyMs,
+          costUSD: reasoningStep.costUSD ?? 0,
+        });
+
+        // STEP 3: Verification
+        const verificationModel =
+          mc?.verification?.model ?? process.env.NEXT_AI_VERIFICATION_MODEL ?? 'gpt-4-turbo';
+        const verificationStep = startStep('verification', verificationModel);
+        const verification = await verifyWithGPT5(reasoning, cites);
+        endStep(verificationStep, {
+          inputTokens: verification.usage.inputTokens,
+          outputTokens: verification.usage.outputTokens,
+        });
+
+        await persistStep(run, verificationStep, {
+          verification_output: {
+            finalAnswer: verification.finalAnswer,
+            flags: verification.flags,
+            trace: verification.trace,
+            confidence: verification.confidence,
+            meta: {
+              model: verificationStep.model,
+              timestamp: new Date().toISOString(),
+            },
+          },
+        });
+        run.steps.push(verificationStep);
+
+        options?.onTrace?.(run.runId, {
+          kind: 'verification',
+          latencyMs: verificationStep.latencyMs,
+          costUSD: verificationStep.costUSD ?? 0,
+        });
+
+        // Finalize run
+        run.flags = verification.flags.map((f: { issue: string }) => f.issue);
+        endRun(run);
+
+        // Persist metadata
+        await persistRunMeta(run);
+        await persistRunIndex(run);
+
+        // Emit completion signal
+        options?.onSignal?.({
+          id: crypto.randomUUID(),
+          type: 'tool.executed',
+          message: `Research completed (confidence: ${verification.confidence}, ${verification.flags.length} flags)`,
+          runId: run.runId,
+          timestamp: new Date().toISOString(),
+          metadata: {
+            runId: run.runId,
+            latencyMs: run.totalLatencyMs,
+            costUSD: run.totalCostUSD,
+            confidence: verification.confidence,
+            flagCount: verification.flags.length,
+          },
+        });
+
+        // Return formatted result
+        return {
+          success: true,
+          answer: verification.finalAnswer,
+          confidence: verification.confidence,
+          citations: cites.map((url, i) => `[${i + 1}] ${url}`).join('\n'),
+          flags:
+            verification.flags.length > 0
+              ? verification.flags
+                  .map(
+                    (f: { claim: string; issue: string; details: string }) =>
+                      `⚠️ ${f.claim}: ${f.issue} (${f.details})`,
+                  )
+                  .join('\n')
+              : 'No verification issues',
+          reasoning: {
+            claims: reasoning.claims.length,
+            assumptions: reasoning.assumptions.length,
+            gaps: reasoning.gaps.length,
+          },
+          trace: {
+            runId: run.runId,
+            latencyMs: run.totalLatencyMs,
+            costUSD: `$${run.totalCostUSD?.toFixed(4)}`,
+            steps: run.steps.map(s => ({
+              kind: s.kind,
+              latencyMs: s.latencyMs,
+              costUSD: `$${s.costUSD?.toFixed(4)}`,
+            })),
+          },
+        };
+      } catch (error) {
+        const errorMessage = error instanceof Error ? error.message : 'Unknown error';
+
+        options?.onSignal?.({
+          id: crypto.randomUUID(),
+          type: 'guardrail.blocked',
+          message: `Research failed: ${errorMessage}`,
+          runId: crypto.randomUUID(),
+          timestamp: new Date().toISOString(),
+          metadata: { error: errorMessage },
+        });
+
+        return {
+          success: false,
+          error: errorMessage,
+          answer: 'Research pipeline failed. Please try again or rephrase your question.',
+        };
+      }
+    },
+  });
+}