refactor llm service architecture with factory pattern (#70)

* refactor llm service architecture with factory pattern

* refactor llm_chat_completion function to improve clarity

* remove docstring for LLMService and chat_completion

* remove unused imports, simplify LLMServiceFactory

Author: Seif-Mamdouh

ai_commit_msg/core/llm_chat_completion.py

@@ -1,25 +1,15 @@
-from ai_commit_msg.services.anthropic_service import AnthropicService
 from ai_commit_msg.services.config_service import ConfigService
-from ai_commit_msg.services.ollama_service import OLlamaService
-from ai_commit_msg.services.openai_service import OpenAiService
+from ai_commit_msg.services.llm_service_factory import LLMServiceFactory
 from ai_commit_msg.utils.logger import Logger
-from ai_commit_msg.utils.models import ANTHROPIC_MODEL_LIST, OPEN_AI_MODEL_LIST
 
 
 def llm_chat_completion(prompt):
     select_model = ConfigService.get_model()
 
-    # TODO - create a factory with a shared interface for calling the LLM models, this will make it easier to add new models
-    ai_gen_commit_msg = None
-    if str(select_model) in OPEN_AI_MODEL_LIST:
-        ai_gen_commit_msg = OpenAiService().chat_with_openai(prompt)
-    elif select_model.startswith("ollama"):
-        ai_gen_commit_msg = OLlamaService().chat_completion(prompt)
-    elif select_model in ANTHROPIC_MODEL_LIST:
-        ai_gen_commit_msg = AnthropicService().chat_completion(prompt)
+    service = LLMServiceFactory.create_service(select_model)
 
-    if ai_gen_commit_msg is None:
+    if service is None:
         Logger().log("Unsupported model: " + select_model)
         return ""
 
-    return ai_gen_commit_msg
+    return service.chat_completion(prompt)

ai_commit_msg/services/anthropic_service.py

@@ -4,9 +4,10 @@
 from ai_commit_msg.utils.models import ANTHROPIC_MODEL_LIST
 from ai_commit_msg.services.config_service import ConfigService
 from ai_commit_msg.utils.error import map_error
+from ai_commit_msg.services.llm_service import LLMService
 
 
-class AnthropicService:
+class AnthropicService(LLMService):
     api_key = ""
 
     def __init__(self):

ai_commit_msg/services/llm_service.py

@@ -0,0 +1,6 @@
+from abc import ABC, abstractmethod
+
+class LLMService(ABC):
+    @abstractmethod
+    def chat_completion(self, messages):
+        pass

ai_commit_msg/services/llm_service_factory.py

@@ -0,0 +1,20 @@
+from ai_commit_msg.services.openai_service import OpenAiService
+from ai_commit_msg.services.anthropic_service import AnthropicService
+from ai_commit_msg.services.ollama_service import OLlamaService
+from ai_commit_msg.utils.models import OPEN_AI_MODEL_LIST, ANTHROPIC_MODEL_LIST
+from ai_commit_msg.utils.logger import Logger
+
+
+class LLMServiceFactory:
+
+    @staticmethod
+    def create_service(model_name):
+        if model_name in OPEN_AI_MODEL_LIST:
+            return OpenAiService()
+        elif model_name.startswith("ollama"):
+            return OLlamaService()
+        elif model_name in ANTHROPIC_MODEL_LIST:
+            return AnthropicService()
+        else:
+            Logger().log(f"Unsupported model: {model_name}")
+            return None

ai_commit_msg/services/ollama_service.py

@@ -2,9 +2,10 @@
 
 from ai_commit_msg.services.config_service import ConfigService
 from ai_commit_msg.utils.logger import Logger
+from ai_commit_msg.services.llm_service import LLMService
 
 
-class OLlamaService:
+class OLlamaService(LLMService):
     def __init__(self):
         self.url = ConfigService.get_ollama_url()
 

ai_commit_msg/services/openai_service.py

@@ -8,9 +8,10 @@
     CONFIG_COLLECTION_KEY,
 )
 from ai_commit_msg.utils.models import OPEN_AI_MODEL_LIST
+from ai_commit_msg.services.llm_service import LLMService
 
 
-class OpenAiService:
+class OpenAiService(LLMService):
     client = None
 
     def __init__(self):
@@ -26,7 +27,7 @@ def __init__(self):
             )
         self.client = OpenAI(api_key=api_key)
 
-    def chat_with_openai(self, messages):
+    def chat_completion(self, messages):
         model_name = ConfigService.get_model()
 
         if model_name not in OPEN_AI_MODEL_LIST:
@@ -41,6 +42,9 @@ def chat_with_openai(self, messages):
         except Exception as e:
             raise map_error("OPENAI", getattr(e, "code", str(e)), e)
 
+    def chat_with_openai(self, messages):
+        return self.chat_completion(messages)
+
     @staticmethod
     def get_openai_api_key():
         raw_json_db = LocalDbService().get_db()[CONFIG_COLLECTION_KEY]