From 9d2be297088c82df4866e907d44d80877a3f546d Mon Sep 17 00:00:00 2001
From: braginini <bangvalo@gmail.com>
Date: Sat, 14 Mar 2026 20:41:54 +0100
Subject: [PATCH] Add Middleware to enrich context

---
 src/client/AIAssistantProvider.tsx |  5 +----
 src/server/handler.ts              | 26 ++++++++++++++++++++++----
 src/server/index.ts                |  2 +-
 src/server/providers/types.ts      | 12 +++++++++++-
 4 files changed, 35 insertions(+), 10 deletions(-)

diff --git a/src/client/AIAssistantProvider.tsx b/src/client/AIAssistantProvider.tsx
index 0885ed8..6abbd7f 100644
--- a/src/client/AIAssistantProvider.tsx
+++ b/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) => {
diff --git a/src/server/handler.ts b/src/server/handler.ts
index f0df66d..3e76b85 100644
--- a/src/server/handler.ts
+++ b/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 };
   }
 
diff --git a/src/server/index.ts b/src/server/index.ts
index 28509d9..d7fe065 100644
--- a/src/server/index.ts
+++ b/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";
\ No newline at end of file
+export type { LLMProvider, ProviderConfig, Middleware } from "./providers/types";
\ No newline at end of file
diff --git a/src/server/providers/types.ts b/src/server/providers/types.ts
index fd18bb0..58f55a3 100644
--- a/src/server/providers/types.ts
+++ b/src/server/providers/types.ts
@@ -7,4 +7,14 @@ export interface LLMProvider {
 export type ProviderConfig = {
   apiKey: string;
   model?: string;
-};
\ No newline at end of file
+};
+
+/**
+ * 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 }>;
\ No newline at end of file