Fixup types and tests

src/helpers.ts

@@ -96,7 +96,9 @@ export function buildMessages(
 /**
  * Build response format object for API from prompt config
  */
-export function buildResponseFormat(promptConfig?: PromptConfig): any {
+export function buildResponseFormat(
+  promptConfig?: PromptConfig
+): { type: 'json_schema'; json_schema: unknown } | undefined {
   if (
     promptConfig?.responseFormat === 'json_schema' &&
     promptConfig.jsonSchema

src/inference.ts

@@ -1,16 +1,30 @@
 import * as core from '@actions/core'
 import ModelClient, { isUnexpected } from '@azure-rest/ai-inference'
 import { AzureKeyCredential } from '@azure/core-auth'
-import { GitHubMCPClient, executeToolCalls } from './mcp.js'
+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?: any // Will contain the processed response format for the API
+  responseFormat?: { type: 'json_schema'; json_schema: unknown } // Processed response format for the API
 }
 
 export interface InferenceResponse {
@@ -41,7 +55,7 @@ export async function simpleInference(
     }
   )
 
-  const requestBody: any = {
+  const requestBody: ChatCompletionsRequestBody = {
     messages: request.messages,
     max_tokens: request.maxTokens,
     model: request.modelName
@@ -84,7 +98,7 @@ export async function mcpInference(
   )
 
   // Start with the pre-processed messages
-  const messages: Array<any> = [...request.messages]
+  const messages: ChatMessage[] = [...request.messages]
 
   let iterationCount = 0
   const maxIterations = 5 // Prevent infinite loops
@@ -93,7 +107,7 @@ export async function mcpInference(
     iterationCount++
     core.info(`MCP inference iteration ${iterationCount}`)
 
-    const requestBody: any = {
+    const requestBody: ChatCompletionsRequestBody = {
       messages: messages,
       max_tokens: request.maxTokens,
       model: request.modelName,

src/main.ts

@@ -8,7 +8,8 @@ import { loadContentFromFileOrInput, buildInferenceRequest } from './helpers.js'
 import {
   loadPromptFile,
   parseTemplateVariables,
-  isPromptYamlFile
+  isPromptYamlFile,
+  PromptConfig
 } from './prompt.js'
 
 const RESPONSE_FILE = 'modelResponse.txt'
@@ -23,7 +24,7 @@ export async function run(): Promise<void> {
     const promptFilePath = core.getInput('prompt-file')
     const inputVariables = core.getInput('input')
 
-    let promptConfig: any = undefined
+    let promptConfig: PromptConfig | undefined = undefined
     let systemPrompt: string | undefined = undefined
     let prompt: string | undefined = undefined