From c2560fc99666f071fb275a3f48774ebfbadaf856 Mon Sep 17 00:00:00 2001
From: Dani Alcala <112832187+clavedeluna@users.noreply.github.com>
Date: Mon, 14 Oct 2024 17:20:30 -0300
Subject: [PATCH] forgot to change top of class typing (#878)

---
 src/codemodder/context.py | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/src/codemodder/context.py b/src/codemodder/context.py
index 0c28e51a..98fa9aac 100644
--- a/src/codemodder/context.py
+++ b/src/codemodder/context.py
@@ -28,6 +28,7 @@
 from codemodder.utils.update_finding_metadata import update_finding_metadata
 
 if TYPE_CHECKING:
+    from azure.ai.inference import ChatCompletionsClient
     from openai import OpenAI
 
     from codemodder.codemods.base_codemod import BaseCodemod
@@ -50,7 +51,8 @@ class CodemodExecutionContext:
     max_workers: int = 1
     tool_result_files_map: dict[str, list[str]]
     semgrep_prefilter_results: ResultSet | None = None
-    llm_client: OpenAI | None = None
+    openai_llm_client: OpenAI | None = None
+    azure_llama_llm_client: ChatCompletionsClient | None = None
 
     def __init__(
         self,