ai-inference/src/inference.ts

import * as core from '@actions/core'
import ModelClient, { isUnexpected } from '@azure-rest/ai-inference'
import { AzureKeyCredential } from '@azure/core-auth'
import { GitHubMCPClient, executeToolCalls, MCPTool, ToolCall } from './mcp.js'
import { handleUnexpectedResponse } from './helpers.js'

interface ChatMessage {
  role: string
  content: string | null
  tool_calls?: ToolCall[]
}

interface ChatCompletionsRequestBody {
  messages: ChatMessage[]
  max_tokens: number
  model: string
  response_format?: { type: 'json_schema'; json_schema: unknown }
  tools?: MCPTool[]
}

export interface InferenceRequest {
  messages: Array<{ role: string; content: string }>
  modelName: string
  maxTokens: number
  endpoint: string
  token: string
  responseFormat?: { type: 'json_schema'; json_schema: unknown } // Processed response format for the API
}

export interface InferenceResponse {
  content: string | null
  toolCalls?: Array<{
    id: string
    type: string
    function: {
      name: string
      arguments: string
    }
  }>
}

/**
 * Simple one-shot inference without tools
 */
export async function simpleInference(
  request: InferenceRequest
): Promise<string | null> {
  core.info('Running simple inference without tools')

  const client = ModelClient(
    request.endpoint,
    new AzureKeyCredential(request.token),
    {
      userAgentOptions: { userAgentPrefix: 'github-actions-ai-inference' }
    }
  )

  const requestBody: ChatCompletionsRequestBody = {
    messages: request.messages,
    max_tokens: request.maxTokens,
    model: request.modelName
  }

  // Add response format if specified
  if (request.responseFormat) {
    requestBody.response_format = request.responseFormat
  }

  const response = await client.path('/chat/completions').post({
    body: requestBody
  })

  if (isUnexpected(response)) {
    handleUnexpectedResponse(response)
  }

  const modelResponse = response.body.choices[0].message.content
  core.info(`Model response: ${modelResponse || 'No response content'}`)

  return modelResponse
}

/**
 * GitHub MCP-enabled inference with tool execution loop
 */
export async function mcpInference(
  request: InferenceRequest,
  githubMcpClient: GitHubMCPClient
): Promise<string | null> {
  core.info('Running GitHub MCP inference with tools')

  const client = ModelClient(
    request.endpoint,
    new AzureKeyCredential(request.token),
    {
      userAgentOptions: { userAgentPrefix: 'github-actions-ai-inference' }
    }
  )

  // Start with the pre-processed messages
  const messages: ChatMessage[] = [...request.messages]

  let iterationCount = 0
  const maxIterations = 5 // Prevent infinite loops

  while (iterationCount < maxIterations) {
    iterationCount++
    core.info(`MCP inference iteration ${iterationCount}`)

    const requestBody: ChatCompletionsRequestBody = {
      messages: messages,
      max_tokens: request.maxTokens,
      model: request.modelName,
      tools: githubMcpClient.tools
    }

    // Add response format if specified (only on first iteration to avoid conflicts)
    if (iterationCount === 1 && request.responseFormat) {
      requestBody.response_format = request.responseFormat
    }

    const response = await client.path('/chat/completions').post({
      body: requestBody
    })

    if (isUnexpected(response)) {
      handleUnexpectedResponse(response)
    }

    const assistantMessage = response.body.choices[0].message
    const modelResponse = assistantMessage.content
    const toolCalls = assistantMessage.tool_calls

    core.info(`Model response: ${modelResponse || 'No response content'}`)

    messages.push({
      role: 'assistant',
      content: modelResponse || '',
      ...(toolCalls && { tool_calls: toolCalls })
    })

    if (!toolCalls || toolCalls.length === 0) {
      core.info('No tool calls requested, ending GitHub MCP inference loop')
      return modelResponse
    }

    core.info(`Model requested ${toolCalls.length} tool calls`)

    // Execute all tool calls via GitHub MCP
    const toolResults = await executeToolCalls(
      githubMcpClient.client,
      toolCalls
    )

    // Add tool results to the conversation
    messages.push(...toolResults)

    core.info('Tool results added, continuing conversation...')
  }

  core.warning(
    `GitHub MCP inference loop exceeded maximum iterations (${maxIterations})`
  )

  // Return the last assistant message content
  const lastAssistantMessage = messages
    .slice()
    .reverse()
    .find((msg) => msg.role === 'assistant')

  return lastAssistantMessage?.content || null
}