From fbabe317a10c5d7b776a84132fd3410d7784897c Mon Sep 17 00:00:00 2001
From: Luke Hinds <luke@stacklok.com>
Date: Mon, 2 Dec 2024 20:27:06 +0000
Subject: [PATCH] Convert async _close_models in LlamaCppInferenceEngine

Fixes RuntimeWarning about unawaited coroutine by converting
the async cleanup method to synchronous since it's called from
del. Also fixes dictionary iteration and adjusts method
name to follow Python protected method conventions.

Resolves: #109
---
 src/codegate/inference/inference_engine.py | 20 ++++++++++----------
 src/codegate/server.py                     |  2 +-
 2 files changed, 11 insertions(+), 11 deletions(-)

diff --git a/src/codegate/inference/inference_engine.py b/src/codegate/inference/inference_engine.py
index f9ae59b1..caaa92a9 100644
--- a/src/codegate/inference/inference_engine.py
+++ b/src/codegate/inference/inference_engine.py
@@ -24,7 +24,16 @@ def __init__(self):
             self.__models = {}
 
     def __del__(self):
-        self.__close_models()
+        self._close_models()
+
+    def _close_models(self):
+        """
+        Closes all open models and samplers
+        """
+        for _, model in self.__models.items():
+            if model._sampler:
+                model._sampler.close()
+            model.close()
 
     async def __get_model(self, model_path, embedding=False, n_ctx=512, n_gpu_layers=0):
         """
@@ -70,12 +79,3 @@ async def embed(self, model_path, content):
         """
         model = await self.__get_model(model_path=model_path, embedding=True)
         return model.embed(content)
-
-    async def __close_models(self):
-        """
-        Closes all open models and samplers
-        """
-        for _, model in self.__models:
-            if model._sampler:
-                model._sampler.close()
-            model.close()
diff --git a/src/codegate/server.py b/src/codegate/server.py
index ec5093ba..3bfa2b42 100644
--- a/src/codegate/server.py
+++ b/src/codegate/server.py
@@ -8,8 +8,8 @@
 from codegate.pipeline.codegate_context_retriever.codegate import CodegateContextRetriever
 from codegate.pipeline.codegate_system_prompt.codegate import CodegateSystemPrompt
 from codegate.pipeline.extract_snippets.extract_snippets import CodeSnippetExtractor
-from codegate.pipeline.secrets.signatures import CodegateSignatures
 from codegate.pipeline.secrets.secrets import CodegateSecrets
+from codegate.pipeline.secrets.signatures import CodegateSignatures
 from codegate.pipeline.version.version import CodegateVersion
 from codegate.providers.anthropic.provider import AnthropicProvider
 from codegate.providers.llamacpp.provider import LlamaCppProvider