Add Middleware to enrich context

src/client/AIAssistantProvider.tsx

@@ -205,11 +205,8 @@ export default function AIAssistantProvider({
 
       if (onExplain) {
         onExplain({ label, query, context: explainCtx, docs: elementDocs, element: explainable });
-        setExplainMode(false);
-        setHoveredEl(null);
-      } else {
-        openChat(query);
       }
+      openChat(query);
     };
 
     const handleKeyDown = (e: KeyboardEvent) => {

src/server/handler.ts

@@ -1,13 +1,14 @@
 import type { Message } from "../types";
 import { AnthropicProvider } from "./providers/anthropic";
 import { OpenAIProvider } from "./providers/openai";
-import type { LLMProvider } from "./providers/types";
+import type { LLMProvider, Middleware } from "./providers/types";
 
 export type AssistantConfig = {
-  provider: "anthropic" | "openai";
-  apiKey: string;
+  provider: "anthropic" | "openai" | LLMProvider;
+  apiKey?: string;
   model?: string;
   systemPrompt?: string;
+  middleware?: Middleware[];
 };
 
 type HandlerOptions = {
@@ -93,6 +94,12 @@ async function parseBody(req: IncomingRequest): Promise<unknown> {
 }
 
 function createProvider(config: AssistantConfig): LLMProvider {
+  if (typeof config.provider === "object") {
+    return config.provider;
+  }
+  if (!config.apiKey) {
+    throw new Error("apiKey is required when using a built-in provider");
+  }
   switch (config.provider) {
     case "anthropic":
       return new AnthropicProvider({ apiKey: config.apiKey, model: config.model });
@@ -105,12 +112,23 @@ function createProvider(config: AssistantConfig): LLMProvider {
 
 export function createAssistant(config: AssistantConfig) {
   const provider = createProvider(config);
+  const middlewares = config.middleware || [];
 
   async function chat(req: ChatRequest): Promise<ChatResponse> {
     if (!req.messages || !Array.isArray(req.messages) || req.messages.length === 0) {
       throw new Error("messages array is required and must not be empty");
     }
-    const reply = await provider.chat(req.messages, config.systemPrompt);
+
+    let messages = req.messages;
+    let systemPrompt = config.systemPrompt;
+
+    for (const mw of middlewares) {
+      const result = await mw(messages, systemPrompt);
+      messages = result.messages;
+      systemPrompt = result.systemPrompt;
+    }
+
+    const reply = await provider.chat(messages, systemPrompt);
     return { reply };
   }
 

src/server/index.ts

@@ -2,4 +2,4 @@ export { createAssistant } from "./handler";
 export type { AssistantConfig } from "./handler";
 export { AnthropicProvider } from "./providers/anthropic";
 export { OpenAIProvider } from "./providers/openai";
-export type { LLMProvider, ProviderConfig } from "./providers/types";
+export type { LLMProvider, ProviderConfig, Middleware } from "./providers/types";

src/server/providers/types.ts

@@ -7,4 +7,14 @@ export interface LLMProvider {
 export type ProviderConfig = {
   apiKey: string;
   model?: string;
-};
+};
+
+/**
+ * Middleware that can transform messages before they reach the LLM.
+ * Useful for RAG context injection, logging, filtering, etc.
+ * Return the (possibly modified) messages array.
+ */
+export type Middleware = (
+  messages: Message[],
+  systemPrompt?: string,
+) => Promise<{ messages: Message[]; systemPrompt?: string }>;