refactor(*): Use a stream for handling the agent communication

src/lib/agent/services/agentService.ts

@@ -1,8 +1,5 @@
 import { DefaultAzureCredential } from '@azure/identity';
-import { AIProjectsClient } from '@azure/ai-projects';
+import { AIProjectsClient, ToolUtility } from '@azure/ai-projects';
-import { ToolUtility } from '@azure/ai-projects';
-import { AI_FOUNDRY_PROJECT_CONNECTION_STRING, AI_MODEL } from '$env/static/private';
-
 import {
   rubricMatcherTool,
   essayAnalyzerTool,
@@ -12,58 +9,64 @@ import {
   selfAssessmentTool,
   spellingAndGrammarCheckerTool
 } from '../tools/index';
+import { AI_FOUNDRY_PROJECT_CONNECTION_STRING, AI_MODEL } from '$env/static/private';
 
-const connectionString = AI_FOUNDRY_PROJECT_CONNECTION_STRING;
+if (!AI_FOUNDRY_PROJECT_CONNECTION_STRING) {
-if (!connectionString) {
+  throw new Error('AI_FOUNDRY_PROJECT_CONNECTION_STRING is not set');
-  throw new Error('AI_FOUNDRY_PROJECT_CONNECTION_STRING is not set in environment variables');
+}
+if (!AI_MODEL) {
+  throw new Error('AI_MODEL is not set');
 }
 
-const aiModel = AI_MODEL || '';
+export const client = AIProjectsClient.fromConnectionString(
-if (!connectionString) {
+  AI_FOUNDRY_PROJECT_CONNECTION_STRING,
-  throw new Error('AI_MODEL is not set in environment variables');
-}
-
-const client = AIProjectsClient.fromConnectionString(
-  connectionString,
   new DefaultAzureCredential()
 );
 
 const codeInterpreterTool = ToolUtility.createCodeInterpreterTool([]);
 
-const tools = [
+const allTools = [
-  codeInterpreterTool.definition,
+  codeInterpreterTool,
-  rubricMatcherTool.definition,
+  rubricMatcherTool,
-  essayAnalyzerTool.definition,
+  essayAnalyzerTool,
-  conceptVerifierTool.definition,
+  conceptVerifierTool,
-  feedbackGeneratorTool.definition,
+  feedbackGeneratorTool,
-  metacognitiveReflectionPromptTool.definition,
+  metacognitiveReflectionPromptTool,
-  selfAssessmentTool.definition,
+  selfAssessmentTool,
-  spellingAndGrammarCheckerTool.definition
+  spellingAndGrammarCheckerTool
 ];
 
-const instructions = "You are a helpful agent that can assist with providing feedback on a student's work for a teacher. You are familiar with all modern pedagogy around summative and formative assessment, and K12 education, and primary and secondary education.";
+const tools = allTools.map(t => t.definition);
 
-const toolResources = {
+const instructions = `
-  ...codeInterpreterTool.resources,
+You are a helpful agent that can assist with providing feedback on a student's work for a teacher.
-  matchRubric: {},
+You are familiar with modern pedagogy around summative and formative assessment in K-12 education.
-  analyzeEssay: {},
+Whenever you need to:
-  conceptVerifier: {},
+ • check the submission against the rubric → call "${rubricMatcherTool.definition.name}"  
-  feedbackGenerator: {},
+ • verify which rubric concepts appear or are missing → call "${conceptVerifierTool.definition.name}"  
-  metacognitiveReflectionPrompt: {},
+ • analyze essay structure → call "${essayAnalyzerTool.definition.name}"  
-  selfAssessment: {},
+ • generate targeted feedback → call "${feedbackGeneratorTool.definition.name}"  
-  spellingAndGrammarChecker: {}
+ • craft a metacognitive prompt → call "${metacognitiveReflectionPromptTool.definition.name}"  
-};
+ • guide self-assessment → call "${selfAssessmentTool.definition.name}"  
+ • check spelling & grammar → call "${spellingAndGrammarCheckerTool.definition.name}"  
+`;
 
-const agent = await client.agents.createAgent(aiModel, {
+const toolResources = allTools.reduce<Record<string, any>>((map, t) => {
+  if (typeof t === 'function') {
+    map[t.definition.name] = t;
+  } else {
+    map[t.definition.name] = t.resources;
+  }
+  return map;
+}, {});
+
+export const agent = await client.agents.createAgent(AI_MODEL, {
   name: `assessment-feedback-agent`,
   instructions,
   temperature: 0.5,
   tools,
   toolResources,
-
   requestOptions: {
-    headers: { "x-ms-enable-preview": "true" },
+    headers: { 'x-ms-enable-preview': 'true' }
-  },
+  }
 });
-
-export { client, agent };

src/lib/agent/tools/spellingAndGrammarCheckerTool.ts

@@ -8,7 +8,7 @@ export const spellingAndGrammarCheckerToolMeta = {
 
 export const spellingAndGrammarCheckerTool = ToolUtility.createFunctionTool({
     name: spellingAndGrammarCheckerToolMeta.key,
-    description: 'Identifies spelling, punctuation, and syntax issues in the submission and suggests corrections.',
+    description: 'Identifies spelling, punctuation, and syntax issues in the submission and suggests corrections. This should result in a clear set of pairs, e.g. "shud" - "should"',
     parameters: {
       type: 'object',
       properties: {

src/lib/utils/ai.ts

@@ -1,10 +1,13 @@
-import { client, agent } from '../agent/services/agentService';
+import { client, agent, toolResources } from '../agent/services/agentService';
-import type { MessageTextContentOutput, ToolOutput } from '@azure/ai-projects';
+import {
+  RunStreamEvent,
+  ErrorEvent,
+  type ThreadRunOutput
+} from '@azure/ai-projects';
 import { PARTYKIT_BASE_URL } from '$env/static/private';
 import type { OpenAIResponse } from './types';
 
-// Prompt template for grading and feedback
+export function buildGradingPrompt(submission: string, task: string): string {
-export function buildGradingPrompt(submission: string, task: string) {
   return `You are an expert secondary school teacher and AI assessment agent. Assess the following student submission in the context of the assignment/task provided.
 
 TASK/ASSIGNMENT:
@@ -14,7 +17,7 @@ STUDENT SUBMISSION:
 ${submission}
 
 Follow these steps:
-1. Grade the work out of 10, using clear, objective criteria.
+1. Grade the work with a letter (A+ is best, E- is worst), using clear, objective criteria.
 2. Identify specific strengths, referencing the success criteria.
 3. Identify misconceptions or areas for improvement, using formative assessment language.
 4. Design an individualized activity or exercise for the student to address their misconceptions or extend their learning. This activity should be:
@@ -28,7 +31,7 @@ Follow these steps:
 RESPONSE FORMAT (respond with a single JSON object, no extra text):
 
 {
-  "grade": "<number or string, e.g. A, B, C+, etc>",
+  "grade": "<number or string, e.g. 8/10, A, B+>",
   "strengths": "<text>",
   "areasForImprovement": "<text>",
   "individualizedActivity": "<text>",
@@ -38,160 +41,152 @@ RESPONSE FORMAT (respond with a single JSON object, no extra text):
   "reasoning": "<step-by-step explanation>"
 }
 
-Respond ONLY with the JSON object, with no preamble or explanation.
+Respond ONLY with the JSON object, with no preamble or explanation.`;
-`;
 }
 
 function fallbackGrade(submission: string, task: string): OpenAIResponse {
-  console.log('using fallback');
+  console.warn('Using fallback grade logic');
   return {
-    grade: '7/10',
+    grade: 'EXAMPLE',
-    strengths: 'Clear argument and good evidence (in response to the task: ' + task + ").",
+    strengths: `Clear argument and good evidence (task: ${task}).`,
-    areasForImprovement: 'Needs deeper analysis and more examples.',
+    areasForImprovement: 'Needs deeper analysis.',
-    individualizedActivity: 'Write a paragraph expanding on your main point and provide two additional examples to support your argument.',
+    individualizedActivity: 'Add two more supporting examples.',
-    reflectionQuestion: 'What was the most challenging part of this assignment for you, and why?',
+    reflectionQuestion: 'Which part was hardest?',
-    teacherSuggestion: 'In the next lesson, review how to develop arguments with supporting evidence and provide a model answer for comparison.',
+    teacherSuggestion: 'Model a paragraph with evidence.',
-    spellingAndGrammar: 'Everything was spelled correctly',
+    spellingAndGrammar: 'No errors detected.',
-    reasoning: 'The submission demonstrates a good structure and evidence, but lacks depth and breadth of analysis. The activity and suggestions are designed to target this gap.'
+    reasoning: 'Standard fallback reasoning.'
   };
 }
 
-export async function gradeSubmissionWithAgent(submission: string, task: string, roomId?: string): Promise<OpenAIResponse> {
+function extractTextFromMessage(msg: { content: any[] }): string {
+  return msg.content
+    .filter(c => c.type === 'text')
+    .map(c => (typeof c.text === 'string' ? c.text : c.text.value ?? ''))
+    .join('\n')
+    .trim();
+}
+
+// Recursively process each run stream, invoke the tools & notify PartyKit
+async function processRunStream(
+  threadId: string,
+  stream: AsyncIterable<RunStreamEvent>,
+  roomId: string
+): Promise<ThreadRunOutput> {
+  for await (const evt of stream) {
+    // if (!evt.event.includes('delta')) {
+    //   console.log('▶️ Stream event:', evt.event);
+    // }
+
+    if (evt.event === RunStreamEvent.ThreadRunRequiresAction) {
+      const runOutput = evt.data as ThreadRunOutput;
+      const calls = runOutput.requiredAction!.submitToolOutputs!.toolCalls!;
+
+      // notify PartyKit
+      await Promise.all(
+        calls.map(async call => {
+          if (!PARTYKIT_BASE_URL) return;
+          const wsCtor = typeof WebSocket !== 'undefined'
+            ? WebSocket
+            : (await import('ws')).default;
+          const ws = new wsCtor(
+            `${PARTYKIT_BASE_URL}/party/tool-usage-server-${roomId}`
+          );
+          await new Promise<void>((res, rej) => {
+            ws.onopen = () => res();
+            ws.onerror = e => rej(e);
+          });
+          ws.send(JSON.stringify({ tool: call.function.name, time: new Date().toISOString() }));
+          ws.close();
+        })
+      );
+
+      // invoke the tools
+      const toolOutputs = await Promise.all(
+        calls.map(async call => {
+          const name = call.function.name;
+          const args = call.arguments!;
           try {
-    if (!client || !agent) throw new Error('Agent not available');
+            const fn = toolResources[name];
+            if (typeof fn !== 'function') {
+              throw new Error(`No tool implementation for "${name}"`);
+            }
+            const result = await fn(args);
+            return { toolCallId: call.id, output: result };
+          } catch (err: any) {
+            return { toolCallId: call.id, output: `Tool error: ${err.message}` };
+          }
+        })
+      );
+
+      // resume the run
+      const resumed = client.agents.submitToolOutputsToRun(
+        threadId,
+        runOutput.id,
+        toolOutputs
+      );
+      const nextStream = await resumed.stream();
+      return processRunStream(threadId, nextStream, roomId);
+    }
+
+    if (evt.event === RunStreamEvent.ThreadRunCompleted) {
+      return evt.data as ThreadRunOutput;
+    }
+
+    if (evt.event === ErrorEvent.Error) {
+      throw new Error(`Agent error: ${JSON.stringify(evt.data)}`);
+    }
+  }
+
+  throw new Error('Stream ended without completion');
+}
+
+export async function gradeSubmissionWithAgent(
+  submission: string,
+  task: string,
+  roomId: string
+): Promise<OpenAIResponse> {
+  if (!client || !agent) {
+    throw new Error('Agent not available');
+  }
 
-    // 1. Create a new thread
   const thread = await client.agents.createThread();
-
-    // 2. Add a message to the thread
   await client.agents.createMessage(thread.id, {
     role: 'user',
     content: buildGradingPrompt(submission, task)
   });
 
-    // 3. Run the agent
+  let initialStream: AsyncIterable<RunStreamEvent>;
-    const run = await client.agents.createRun(thread.id, agent.id, {});
-
-    // 4. Poll for run completion
-    let runResult;
-    for (let i = 0; i < 60; i++) { // up to ~60 seconds
-      runResult = await client.agents.getRun(thread.id, run.id);
-      console.log('Agent run status:', runResult.status);
-
-      if (runResult.status === 'completed') break;
-      if (runResult.status === 'failed') {
-        const messages = await client.agents.listMessages(thread.id);
-        const agentMessage = messages.data.find((m: { role: string; content: any[] }) => m.role === 'agent');
-        let errorText = '';
-        if (agentMessage && Array.isArray(agentMessage.content) && agentMessage.content.length > 0) {
-          const contentItem = agentMessage.content[0];
-          if (contentItem.type === 'text' && contentItem.text && typeof contentItem.text.value === 'string') {
-            errorText = contentItem.text.value;
-          } else if (typeof contentItem.value === 'string') {
-            errorText = contentItem.value;
-          }
-        }
-        throw new Error(`Agent run failed. Status: failed. Message: ${errorText}`);
-      }
-
-      // Handle tool calls if agent requires action
-      if (runResult.status === 'requires_action' && runResult.requiredAction?.submitToolOutputs.toolCalls) {
-        const toolCalls = runResult.requiredAction.submitToolOutputs.toolCalls;
-        const toolOutputs: ToolOutput[] = await Promise.all(toolCalls.map(async (call: any) => {
-          let output = '';
-          const args = JSON.parse(call.function.arguments);
-
-          if (roomId) {
   try {
-              const ws = new (typeof WebSocket !== 'undefined' ? WebSocket : (await import('ws')).default)(`${PARTYKIT_BASE_URL}/party/tool-usage-server-${roomId}`);
+    const runInvoker = client.agents.createRun(thread.id, agent.id, {
-
+      parallelToolCalls: false
-              // Wait for connection to open
-              await new Promise<void>((resolve, reject) => {
-                ws.onopen = () => resolve();
-                ws.onerror = (err) => reject(err);
     });
-
+    initialStream = await runInvoker.stream();
-              ws.send(JSON.stringify({ tool: call.function.name, time: new Date().toISOString() }));
-              
-              ws.close();
   } catch (err) {
-              console.error('Failed to send tool usage event to PartyKit:', err);
+    console.error('Failed to start run:', err);
-            }
-          }
-
-          if (
-            call.function.name === 'matchRubric' ||
-            call.function.name === 'analyzeEssay'
-          ) {
-            // LLM will handle the tool logic, just submit empty output
-            output = '';
-          } else {
-            output = `Unknown tool: ${call.function.name}`;
-          }
-
-          return {
-            toolCallId: call.id,
-            output: typeof output === 'string' ? output : JSON.stringify(output)
-          };
-        }));
-        await client.agents.submitToolOutputsToRun(thread.id, run.id, toolOutputs);
-      }
-
-      await new Promise(res => setTimeout(res, 1000));
-    }
-    if (!runResult || runResult.status !== 'completed') {
-      const messages = await client.agents.listMessages(thread.id);
-      const agentMessage = messages.data.find((m: { role: string; content: any[] }) => m.role === 'agent');
-      let errorText = '';
-      if (agentMessage && Array.isArray(agentMessage.content) && agentMessage.content.length > 0) {
-        const contentItem = agentMessage.content[0];
-        if (contentItem.type === 'text' && contentItem.text && typeof contentItem.text.value === 'string') {
-          errorText = contentItem.text.value;
-        } else if (typeof contentItem.value === 'string') {
-          errorText = contentItem.value;
-        }
-      }
-      throw new Error(`Agent run did not complete in time. Last status: ${runResult?.status}. Last message: ${errorText}`);
-    }
-
-    // 5. Get agent's feedback message
-    const messages = await client.agents.listMessages(thread.id);
-    const agentMessage = messages.data.find((m: { role: string; content: any[] }) => m.role === 'assistant');
-    // Safely extract text content from agent message
-    let text = '';
-    if (agentMessage && Array.isArray(agentMessage.content) && agentMessage.content.length > 0) {
-      const contentItem = agentMessage.content[0];
-      if (contentItem.type === 'text' && 'text' in contentItem) {
-        text = (contentItem as MessageTextContentOutput).text.value;
-      } else if ('value' in contentItem && typeof (contentItem as any).value === 'string') {
-        text = (contentItem as any).value;
-      }
-    }
-
-    // 6. Parse agent's response into OpenAIResponse format
-    let parsed: OpenAIResponse | null = null;
-    try {
-      parsed = JSON.parse(text);
-    } catch (err) {
-      console.error(err);
-    }
-    if (parsed && typeof parsed === 'object' && parsed.grade) {
-      return parsed;
-    }
-    // If parsing fails, return raw text in reasoning
-    return {
-      grade: '',
-      strengths: '',
-      areasForImprovement: '',
-      individualizedActivity: '',
-      reflectionQuestion: '',
-      teacherSuggestion: '',
-      spellingAndGrammar: '',
-      reasoning: text || 'No feedback generated.'
-    };
-  } catch (err) {
-    console.error(err);
-    // Fallback to local/demo logic
     return fallbackGrade(submission, task);
   }
+
+  // process all tool calls
+  let finalRun: ThreadRunOutput;
+  try {
+    finalRun = await processRunStream(thread.id, initialStream, roomId);
+  } catch (err) {
+    console.error('Agent streaming failed:', err);
+    return fallbackGrade(submission, task);
+  }
+
+  const msgs = await client.agents.listMessages(thread.id);
+  const assistant = msgs.data.find(m => m.role === 'assistant');
+  const raw = assistant ? extractTextFromMessage(assistant) : '';
+
+  // strip anything before the JSON object
+  const match = raw.match(/\{[\s\S]*\}$/);
+  const jsonText = match ? match[0] : raw;
+
+  try {
+    return JSON.parse(jsonText) as OpenAIResponse;
+  } catch (parseErr) {
+    console.error('JSON parse failed, returning fallback:', parseErr);
+    return { ...fallbackGrade(submission, task), reasoning: raw };
+  }
 }