File size: 8,577 Bytes
6cdce85
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
import { NextRequest } from 'next/server';
import { createVLMClient } from '@/lib/api/vlm-client';
import { ALLOWED_TOPICS, BLOCKED_INPUT_PATTERNS } from '@/config/constants';

export const maxDuration = 120;

interface MessageContent {
  type: 'text' | 'image_url';
  text?: string;
  image_url?: { url: string };
}

interface ChatMessage {
  role: 'system' | 'user' | 'assistant';
  content: string | MessageContent[];
}

interface ChatRequestBody {
  messages: ChatMessage[];
  stream?: boolean;
}

interface ContentValidation {
  valid: boolean;
  reason?: string;
  isOffTopic?: boolean;
}

/**
 * Extract text content from a message for validation
 */
function extractTextContent(content: string | MessageContent[]): string {
  if (typeof content === 'string') {
    return content;
  }
  return content
    .filter((c): c is MessageContent & { type: 'text'; text: string } => c.type === 'text' && !!c.text)
    .map(c => c.text)
    .join(' ');
}

/**
 * Validate user input for malicious patterns and topic relevance
 */
function validateUserInput(text: string): ContentValidation {
  const lowerText = text.toLowerCase();
  
  // Check for blocked patterns (prompt injection, harmful content, etc.)
  for (const pattern of BLOCKED_INPUT_PATTERNS) {
    if (pattern.test(text)) {
      return {
        valid: false,
        reason: "I can't process this request. Please ask a question related to quantum computing, Qiskit, physics, or mathematics.",
      };
    }
  }
  
  // Check message length (prevent abuse)
  if (text.length > 10000) {
    return {
      valid: false,
      reason: 'Message too long. Please keep your question under 10,000 characters.',
    };
  }
  
  // Check if the message contains any relevant topic keywords
  // Images are always allowed (circuit diagrams, Bloch spheres, etc.)
  const hasImage = text.includes('[IMAGE]') || text.length < 20; // Short messages might be follow-ups
  
  if (!hasImage) {
    const words = lowerText.split(/\s+/);
    const hasRelevantTopic = ALLOWED_TOPICS.some(topic => {
      // Check for whole word or part of compound word
      return words.some(word => 
        word.includes(topic.toLowerCase()) || 
        topic.toLowerCase().includes(word)
      );
    });
    
    // Also check for common question patterns
    const isQuestion = /^(what|how|why|when|where|can|could|would|should|is|are|do|does|explain|describe|help|show|create|implement|write|generate|build|make)/i.test(lowerText.trim());
    const hasCodeContext = /```|def\s|import\s|class\s|function|circuit/i.test(text);
    
    // Be permissive: if it's a question or has code context, allow it
    // The model will redirect off-topic questions anyway
    if (!hasRelevantTopic && !isQuestion && !hasCodeContext && text.length > 50) {
      return {
        valid: true, // Still valid, but flag as potentially off-topic
        isOffTopic: true,
      };
    }
  }
  
  return { valid: true };
}

/**
 * Create off-topic response message
 */
function createOffTopicResponse(): string {
  return `I'm **Quantum Assistant**, specialized in quantum computing, Qiskit, physics, and related mathematics.

I can help you with:
- 🔬 **Quantum Computing**: Circuits, gates, algorithms, error correction
- 💻 **Qiskit**: Code generation, debugging, best practices
- 📐 **Physics & Math**: Quantum mechanics, linear algebra, probability
- 🤖 **Quantum ML**: Variational algorithms, optimization, hybrid systems

**Please ask a question related to these topics!**

For example:
- "How do I create a Bell state in Qiskit?"
- "Explain the Grover's algorithm"
- "What is quantum entanglement?"`;
}

function isConnectionError(error: unknown): boolean {
  if (error instanceof Error) {
    const message = error.message.toLowerCase();
    const cause = (error as Error & { cause?: Error })?.cause;

    if (message.includes('fetch failed') || message.includes('econnrefused')) {
      return true;
    }

    if (cause && 'code' in cause && cause.code === 'ECONNREFUSED') {
      return true;
    }
  }
  return false;
}

function createErrorMessage(isConnection: boolean): string {
  if (isConnection) {
    const modelUrl = process.env.DEMO_MODEL_URL || 'http://localhost:8000/v1';
    return `**Model Server Not Available**\n\nCould not connect to the model at:\n\`${modelUrl}\`\n\n**To use the chat feature:**\n1. Start a VLM server (vLLM, Ollama, etc.)\n2. Configure \`.env.local\` with your endpoint:\n\`\`\`\nDEMO_MODEL_URL=http://your-server:port/v1\nDEMO_MODEL_NAME=your-model-name\nDEMO_API_KEY=your-api-key\n\`\`\`\n3. Restart the demo server\n\n*Examples panel still works - try selecting a test sample!*`;
  }
  return 'An error occurred while processing your request.';
}

export async function POST(request: NextRequest) {
  try {
    const body: ChatRequestBody = await request.json();
    const { messages, stream = true } = body;

    if (!messages || !Array.isArray(messages) || messages.length === 0) {
      return new Response(
        JSON.stringify({ error: 'Invalid request: messages array required' }),
        { status: 400, headers: { 'Content-Type': 'application/json' } }
      );
    }

    // Find the last user message for validation
    const userMessages = messages.filter(m => m.role === 'user');
    const lastUserMessage = userMessages[userMessages.length - 1];
    
    if (lastUserMessage) {
      const userText = extractTextContent(lastUserMessage.content);
      const validation = validateUserInput(userText);
      
      // If input is invalid (malicious/harmful), return error
      if (!validation.valid && validation.reason) {
        const encoder = new TextEncoder();
        
        if (stream) {
          const errorStream = new ReadableStream({
            start(controller) {
              const data = JSON.stringify({ content: validation.reason, done: false });
              controller.enqueue(encoder.encode(`data: ${data}\n\n`));
              controller.enqueue(encoder.encode(`data: ${JSON.stringify({ done: true })}\n\n`));
              controller.close();
            },
          });
          
          return new Response(errorStream, {
            headers: {
              'Content-Type': 'text/event-stream',
              'Cache-Control': 'no-cache',
              'Connection': 'keep-alive',
            },
          });
        } else {
          return new Response(
            JSON.stringify({ content: validation.reason }),
            { headers: { 'Content-Type': 'application/json' } }
          );
        }
      }
    }

    const client = createVLMClient();

    if (stream) {
      const encoder = new TextEncoder();

      const readableStream = new ReadableStream({
        async start(controller) {
          try {
            for await (const chunk of client.chatStream(messages)) {
              const data = JSON.stringify({ content: chunk, done: false });
              controller.enqueue(encoder.encode(`data: ${data}\n\n`));
            }
            controller.enqueue(encoder.encode(`data: ${JSON.stringify({ done: true })}\n\n`));
            controller.close();
          } catch (error) {
            console.error('Stream error:', error);
            const isConnection = isConnectionError(error);
            const errorMessage = isConnection
              ? createErrorMessage(true)
              : (error instanceof Error ? error.message : 'Stream error occurred');

            controller.enqueue(
              encoder.encode(`data: ${JSON.stringify({ error: errorMessage, done: true })}\n\n`)
            );
            controller.close();
          }
        },
      });

      return new Response(readableStream, {
        headers: {
          'Content-Type': 'text/event-stream',
          'Cache-Control': 'no-cache',
          'Connection': 'keep-alive',
        },
      });
    } else {
      const response = await client.chat(messages);
      return new Response(
        JSON.stringify({ content: response }),
        { headers: { 'Content-Type': 'application/json' } }
      );
    }
  } catch (error) {
    console.error('Chat API error:', error);

    if (isConnectionError(error)) {
      return new Response(
        JSON.stringify({ error: createErrorMessage(true) }),
        { status: 503, headers: { 'Content-Type': 'application/json' } }
      );
    }

    const errorMessage =
      error instanceof Error ? error.message : 'Internal server error';

    return new Response(
      JSON.stringify({ error: errorMessage }),
      { status: 500, headers: { 'Content-Type': 'application/json' } }
    );
  }
}