USTORY-549: Provide model name command line argument

.gitignore

@@ -172,3 +172,7 @@ cython_debug/
 
 # macOS system files
 .DS_Store
+
+# output
+stdout
+output_images/

ai_feedback/__main__.py

@@ -1,5 +1,4 @@
 import argparse
-import json
 import os
 import os.path
 import sys
@@ -9,6 +8,7 @@
 from . import code_processing, image_processing, text_processing
 from .helpers import arg_options
 from .helpers.constants import HELP_MESSAGES
+from .models import ModelFactory
 
 _TYPE_BY_EXTENSION = {
     '.c': 'C',
@@ -192,17 +192,17 @@ def main() -> int:
     parser.add_argument("--solution", type=str, required=False, default="", help=HELP_MESSAGES["solution"])
     parser.add_argument("--question", type=str, required=False, help=HELP_MESSAGES["question"])
     parser.add_argument(
-        "--model",
+        "--provider",
         type=str,
-        choices=arg_options.get_enum_values(arg_options.Models),
+        choices=ModelFactory.get_available_providers(),
         required=True,
-        help=HELP_MESSAGES["model"],
+        help=HELP_MESSAGES["provider"],
     )
     parser.add_argument(
-        "--remote_model",
+        "--model_name",
         type=str,
         required=False,
-        help=HELP_MESSAGES["remote_model"],
+        help=HELP_MESSAGES["model_name"],
     )
     parser.add_argument(
         "--output",
@@ -294,23 +294,32 @@ def main() -> int:
     if args.prompt_text:
         prompt_content += args.prompt_text
 
+    try:
+        model_args = {'model_name': args.model_name} if args.model_name else {}
+        model = ModelFactory.create(args.provider, **model_args)
+    except ValueError as e:
+        print(f"Error: {e}")
+        sys.exit(1)
+
     if args.scope == "image":
         prompt = {"prompt_content": prompt_content}
-        request, response = image_processing.process_image(args, prompt, system_instructions, marking_instructions)
+        request, response = image_processing.process_image(
+            model, args, prompt, system_instructions, marking_instructions
+        )
     elif args.scope == "text":
         request, response = text_processing.process_text(
-            args, prompt_content, system_instructions, marking_instructions
+            model, args, prompt_content, system_instructions, marking_instructions
         )
     else:
         request, response = code_processing.process_code(
-            args, prompt_content, system_instructions, marking_instructions
+            model, args, prompt_content, system_instructions, marking_instructions
         )
 
     markdown_template = load_markdown_template(args.output_template)
     timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
     output_text = markdown_template.format(
         question=args.question or "N/A",
-        model=args.model,
+        model=args.provider,
         request=request,
         response=response,
         timestamp=timestamp,

ai_feedback/code_processing.py

@@ -1,17 +1,15 @@
 import os
-import sys
 from pathlib import Path
 from typing import Callable, Optional, Tuple
 
-from .helpers.arg_options import model_mapping
 from .helpers.file_converter import rename_files
 from .helpers.template_utils import render_prompt_template
 
 EXPECTED_SUFFIXES = ["_solution", "test_output", "_submission"]
 
 
 def process_code(
-    args, prompt: str, system_instructions: str, marking_instructions: Optional[str] = None
+    model, args, prompt: str, system_instructions: str, marking_instructions: Optional[str] = None
 ) -> Tuple[str, str]:
     """
     Processes assignment files and generates a response using the selected model.
@@ -68,16 +66,6 @@ def process_code(
         marking_instructions=marking_instructions,
     )
 
-    if args.model in model_mapping:
-        model_class = model_mapping[args.model]
-        if model_class.__name__ == 'RemoteModel' and args.remote_model:
-            model = model_class(model_name=args.remote_model)
-        else:
-            model = model_class()
-    else:
-        print("Invalid model selected for code scope.")
-        sys.exit(1)
-
     if args.scope == "code":
         if args.question:
             request, response = model.generate_response(

ai_feedback/helpers/arg_options.py

@@ -1,7 +1,5 @@
 from enum import Enum
 
-from .. import models
-
 
 def get_enum_values(enum_class: type[Enum]) -> list[str]:
     """
@@ -68,17 +66,6 @@ def __str__(self):
         return self.value
 
 
-model_mapping = {
-    "deepSeek-R1:70B": models.DeepSeekModel,
-    "openai": models.OpenAIModel,
-    "openai-vector": models.OpenAIModelVector,
-    "codellama:latest": models.CodeLlamaModel,
-    "claude-3.7-sonnet": models.ClaudeModel,
-    "remote": models.RemoteModel,
-    "deepSeek-v3": models.DeepSeekV3Model,
-}
-
-
 class Models(Enum):
     """
     Enum representing the available AI model types.

ai_feedback/helpers/constants.py

@@ -7,8 +7,8 @@
     "submission": "The file path for the submission file.",
     "solution": "The file path for the solution file.",
     "question": "The specific question number to analyze within the assignment (if applicable).",
-    "model": "The name of the LLM model to use for evaluation.",
-    "remote_model": "When using --remote=model, this option specifies the remote model to use.",
+    "provider": "The name of the LLM provider to use for evaluation.",
+    "model_name": "The name of the LLM model to use for evaluation, this option specifies the model to use.",
     "output": "Format to display the output response.",
     "llama_mode": "Specifies how to invoke llama.cpp: either directly via its command‐line interface (CLI) or by sending requests to a running llama-server instance.",
     "test_output": "The output of tests from evaluating the assignment.",

ai_feedback/helpers/image_extractor.py

@@ -8,6 +8,12 @@
 from typing import Any, Dict, List, Optional
 
 
+def encode_image(image_path: os.PathLike) -> bytes:
+    """Encodes the image found at {image_path} to a base64 string"""
+    with open(image_path, "rb") as image_file:
+        return base64.b64encode(image_file.read()).decode("utf-8")
+
+
 def extract_images(input_notebook_path: os.PathLike, output_directory: os.PathLike, output_name: str) -> List[Path]:
     image_paths = []
     with open(input_notebook_path, "r") as file:

ai_feedback/image_processing.py

@@ -1,95 +1,17 @@
-import base64
+import sys
 from pathlib import Path
 from typing import Optional
 
-from anthropic import Anthropic
-from dotenv import load_dotenv
-from ollama import Image, Message, chat
-from openai import OpenAI
+from ollama import Image, Message
 from PIL import Image as PILImage
 
-from .helpers.arg_options import Models
 from .helpers.image_extractor import extract_images, extract_qmd_python_images
 from .helpers.image_reader import *
 from .helpers.template_utils import render_prompt_template
-from .models.RemoteModel import RemoteModel
-
-
-def encode_image(image_path: os.PathLike) -> bytes:
-    """Encodes the image found at {image_path} to a base64 string"""
-    with open(image_path, "rb") as image_file:
-        return base64.b64encode(image_file.read()).decode("utf-8")
-
-
-def openai_call(message: Message, model: str) -> str | None:
-    """Sends a request to OpenAI"""
-    # Load environment variables from .env file
-    load_dotenv()
-    client = OpenAI()
-    images = [
-        {
-            "type": "image_url",
-            "image_url": {"url": f"data:image/png;base64,{encode_image(image.value)}"},
-        }
-        for image in message.images
-    ]
-    completion = client.chat.completions.create(
-        model=model,
-        messages=[
-            {
-                "role": "user",
-                "content": [
-                    {
-                        "type": "text",
-                        "text": message.content,
-                    }
-                ]
-                + images,
-            }
-        ],
-        temperature=0.33,
-    )
-    return completion.choices[0].message.content
-
-
-def anthropic_call(message: Message, model: str) -> str | None:
-    """Sends a request to OpenAI"""
-    # Load environment variables from .env file
-    load_dotenv()
-    client = Anthropic(api_key=os.getenv("CLAUDE_API_KEY"))
-    images = [
-        {
-            "type": "image",
-            "source": {
-                "type": "base64",
-                "media_type": "image/png",
-                "data": f"{encode_image(image.value)}",
-            },
-        }
-        for image in message.images
-    ]
-    message = client.messages.create(
-        max_tokens=2048,
-        model=model,
-        messages=[
-            {
-                "role": "user",
-                "content": [
-                    {
-                        "type": "text",
-                        "text": message.content,
-                    }
-                ]
-                + images,
-            }
-        ],
-        temperature=0.33,
-    )
-    return message.content[0].text
 
 
 def process_image(
-    args, prompt: dict, system_instructions: str, marking_instructions: Optional[str] = None
+    model, args, prompt: dict, system_instructions: str, marking_instructions: Optional[str] = None
 ) -> tuple[str, str]:
     """Generates feedback for an image submission.
     Returns the LLM prompt delivered and the returned response."""
@@ -159,27 +81,9 @@ def process_image(
 
         # Prompt the LLM
         requests.append(f"{message.content}\n\n{[str(image.value) for image in message.images]}")
-        if args.model == Models.OPENAI.value:
-            responses.append(openai_call(message, model="gpt-4o"))
-        elif args.model == Models.CLAUDE.value:
-            responses.append(anthropic_call(message, model="claude-3-7-sonnet-20250219"))
-        elif args.model == Models.REMOTE.value:
-            if args.remote_model:
-                model = RemoteModel(model_name=args.remote_model)
-            else:
-                model = RemoteModel()
 
-            _request, response = model.generate_response(
-                rendered_prompt,
-                args.submission,
-                system_instructions=system_instructions,
-                question=question,
-                submission_image=args.submission_image,
-                json_schema=args.json_schema,
-                model_options=args.model_options,
-            )
-            responses.append(str(response))
-        else:
-            responses.append(chat(model=args.model, messages=[message], options={"temperature": 0.33}).message.content)
+        args.rendered_prompt = rendered_prompt
+        args.system_instructions = system_instructions
+        responses.append(model.process_image(message, args))
 
     return "\n\n---\n\n".join(requests), "\n\n---\n\n".join(responses)

ai_feedback/models/ClaudeModel.py

@@ -4,7 +4,9 @@
 
 import anthropic
 from dotenv import load_dotenv
+from ollama import Message
 
+from ..helpers.image_extractor import encode_image
 from ..helpers.model_options_helpers import cast_to_type, claude_option_schema
 from .Model import Model
 
@@ -13,11 +15,12 @@
 
 
 class ClaudeModel(Model):
-    def __init__(self) -> None:
+    def __init__(self, model_name: str = None) -> None:
         """
         Initializes the ClaudeModel with the Anthropic client using an API key.
         """
-        super().__init__()
+        super().__init__(model_name)
+        self.model_name = model_name if model_name else "claude-3-7-sonnet-20250219"
         self.client = anthropic.Anthropic(api_key=os.getenv("CLAUDE_API_KEY"))
 
     def generate_response(
@@ -62,7 +65,7 @@ def generate_response(
 
         # Construct request parameters
         request_kwargs = {
-            "model": "claude-3-7-sonnet-20250219",
+            "model": self.model_name,
             "system": system_instructions,
             "messages": [{"role": "user", "content": request}],
             **model_options,
@@ -75,3 +78,35 @@ def generate_response(
             return None
 
         return prompt, response.content[0].text
+
+    def process_image(self, message: Message, args) -> str:
+        """Sends a request to Claude"""
+        images = [
+            {
+                "type": "image",
+                "source": {
+                    "type": "base64",
+                    "media_type": "image/png",
+                    "data": f"{encode_image(image.value)}",
+                },
+            }
+            for image in message.images
+        ]
+        response = self.client.messages.create(
+            max_tokens=2048,
+            model=self.model_name,
+            messages=[
+                {
+                    "role": "user",
+                    "content": [
+                        {
+                            "type": "text",
+                            "text": message.content,
+                        }
+                    ]
+                    + images,
+                }
+            ],
+            temperature=0.33,
+        )
+        return response.content[0].text

ai_feedback/models/CodeLlamaModel.py

@@ -11,13 +11,12 @@
 
 class CodeLlamaModel(Model):
 
-    def __init__(self) -> None:
+    def __init__(self, model_name: str = None) -> None:
         """
         Initializes the CodeLlamaModel with configuration for the model and system instructions.
         """
-        self.model = {
-            "model": "codellama:latest",
-        }
+        super().__init__(model_name)
+        self.model_name = model_name if model_name else "codellama:latest"
 
     def generate_response(
         self,
@@ -64,7 +63,7 @@ def generate_response(
         model_options = cast_to_type(ollama_option_schema, model_options)
 
         response = ollama.chat(
-            model=self.model["model"],
+            model=self.model_name,
             messages=[
                 {"role": "system", "content": system_instructions},
                 {"role": "user", "content": prompt},