feat: allow to introduce any model for supported agents (#54)

sdn4z · web-flow · commit 961743628144 · 2025-07-22T17:50:43.000+02:00
diff --git a/eval/cli.py b/eval/cli.py
@@ -26,6 +26,7 @@
     ),
 )
 @click.option("--prompt", type=str, help=("Which prompt to use."))
+@click.option("--model", type=str, default=None, help=("Which prompt to use."))
 @click.option(
     "--prompt-file",
     type=str,
@@ -55,6 +56,7 @@ def run(
     prompt_file: str,
     env_file: str,
     tag: str | None = None,
+    model: str | None = None,
 ) -> None:
     load_dotenv(env_file)
     config_from_file = EvalFileConfig.get_config_from_file(config_section=config, path=config_file)
@@ -65,6 +67,7 @@ def run(
             "prompt": prompt or config_from_file.prompt,
             "score_threshold": score or config_from_file.score_threshold,
             "samples": samples or config_from_file.samples,
+            "model": model or config_from_file.model,
         }
     )
 
@@ -76,6 +79,7 @@ def run(
         tag=tag,
         agent=eval_config.agent,
         prompt=configured_prompts.get_prompt(eval_config.prompt),
+        model=eval_config.model,
     )
 
 
diff --git a/eval/evaluator.py b/eval/evaluator.py
@@ -1,6 +1,6 @@
 import logging
 
-from lightman_ai.ai.utils import get_agent_instance_from_agent_name
+from lightman_ai.ai.utils import get_agent_class_from_agent_name
 from lightman_ai.article.models import Article
 
 from eval.classifier import Classifier
@@ -11,16 +11,19 @@
 
 
 def eval(
+    *,
     prompt: str,
     score_threshold: int,
     relevant_articles: set[Article],
     non_relevant_articles: set[Article],
     samples: int,
     tag: str | None,
     agent: str,
+    model: str | None,
 ) -> None:
+    agent_class = get_agent_class_from_agent_name(agent)
     classified_articles = Classifier(
-        agent=get_agent_instance_from_agent_name(agent)(prompt),
+        agent=agent_class(prompt, model=model),
         score=score_threshold,
         relevant_articles=relevant_articles,
         non_relevant_articles=non_relevant_articles,
@@ -36,6 +39,7 @@ def eval(
         prompt=prompt,
         score=score_threshold,
         logger=logger,
+        model=model or agent_class._default_model_name,
     )
 
     results_template.save()
diff --git a/eval/templates.py b/eval/templates.py
@@ -17,6 +17,7 @@ def __init__(
         samples: int,
         prompt: str,
         score: int,
+        model: str,
         logger: logging.Logger | None = None,
     ) -> None:
         self.results_metrics = results_metrics
@@ -26,6 +27,7 @@ def __init__(
         self.prompt = prompt
         self.score = score
         self.results_dir = Path(RESULTS_DIR)
+        self.model = model
         self.logger = logger or logging.getLogger("eval")
 
     @property
@@ -55,6 +57,7 @@ def _get_summary(self) -> str:
 # Summary
 - Tag: {self.tag or "-"}
 - Agent: {self.agent}
+- Model: {self.model}
 - Samples: {self.samples}
 - Score threshold: {self.score}
 - Prompt: \n {self.prompt}
diff --git a/justfile b/justfile
@@ -1,11 +1,13 @@
 # VARIABLE DEFINITIONS
 venv := ".venv"
 python_version :="3.13"
+run := "poetry run"
+eval_path := "eval/cli.py"
 
 venv-exists := path_exists(venv)
 
 
-target_dirs := "src tests"
+target_dirs := "src tests eval"
 
 # ALIASES
 alias t := test
@@ -19,6 +21,9 @@ venv:
     uv sync --frozen --all-extras; \
     fi
 
+# Runs the evaluation script
+eval *args: venv
+    PYTHONPATH=. {{ run }} python {{ eval_path }} {{ args }}
 
 # Cleans all artifacts generated while running this project, including the virtualenv.
 clean: 
@@ -27,17 +32,17 @@ clean:
 
 # Runs the tests with the specified arguments (any path or pytest argument).
 test *test-args='': venv
-    uv run pytest {{ test-args }} --no-cov 
+    {{ run }}  pytest {{ test-args }} --no-cov 
 
 # Runs all tests including coverage report.
 test-all: venv
-    uv run pytest
+    {{ run }}  pytest
 
 # Format all code in the project.
 format:  venv
-    uv run ruff check {{ target_dirs }} --fix
+    {{ run }} ruff check {{ target_dirs }}
 
 # Lint all code in the project.
 lint: venv
-    uv run ruff check {{ target_dirs }}
-    uv run mypy {{ target_dirs }}
+    {{ run }}  ruff check {{ target_dirs }}
+    {{ run }}  mypy {{ target_dirs }}
diff --git a/lightman.toml b/lightman.toml
@@ -2,7 +2,6 @@
 agent = 'openai'
 score_threshold = 8
 prompt = 'classify'
-
 service_desk_project_key=''
 service_desk_request_id_type=''
 
@@ -11,6 +10,7 @@ agent = 'openai'
 score_threshold = 8
 prompt = 'classify'
 samples = 3
+model = 'gpt-4o'
 
 [prompts]
 classify = ""
diff --git a/src/lightman_ai/ai/base/agent.py b/src/lightman_ai/ai/base/agent.py
@@ -9,13 +9,13 @@
 
 
 class BaseAgent(ABC):
-    _model: type[OpenAIModel] | type[GoogleModel]
-    _model_name: str
+    _class: type[OpenAIModel] | type[GoogleModel]
+    _default_model_name: str
 
-    def __init__(self, system_prompt: str, logger: logging.Logger | None = None) -> None:
-        model = self._model(self._model_name)
+    def __init__(self, system_prompt: str, model: str | None = None, logger: logging.Logger | None = None) -> None:
+        agent_model = self._class(model or self._default_model_name)
         self.agent: Agent[Never, SelectedArticlesList] = Agent(
-            model=model, output_type=SelectedArticlesList, system_prompt=system_prompt
+            model=agent_model, output_type=SelectedArticlesList, system_prompt=system_prompt
         )
         self.logger = logger or logging.getLogger("lightman")
 
diff --git a/src/lightman_ai/ai/gemini/agent.py b/src/lightman_ai/ai/gemini/agent.py
@@ -9,8 +9,8 @@
 class GeminiAgent(BaseAgent):
     """Class that provides an interface to operate with the Gemini model."""
 
-    _model = GoogleModel
-    _model_name = "gemini-2.5-pro-preview-05-06"
+    _class = GoogleModel
+    _default_model_name = "gemini-2.5-pro-preview-05-06"
 
     @override
     def _run_prompt(self, prompt: str) -> SelectedArticlesList:
diff --git a/src/lightman_ai/ai/openai/agent.py b/src/lightman_ai/ai/openai/agent.py
@@ -11,8 +11,8 @@
 class OpenAIAgent(BaseAgent):
     """Class that provides an interface to operate with the OpenAI model."""
 
-    _model = OpenAIModel
-    _model_name = "gpt-4.1"
+    _class = OpenAIModel
+    _default_model_name = "gpt-4.1"
 
     def _execute_agent(self, prompt: str) -> AgentRunResult[SelectedArticlesList]:
         with map_openai_exceptions():
diff --git a/src/lightman_ai/ai/utils.py b/src/lightman_ai/ai/utils.py
@@ -5,7 +5,7 @@
 AGENT_MAPPING = {"openai": OpenAIAgent, "gemini": GeminiAgent}
 
 
-def get_agent_instance_from_agent_name(agent: str) -> type[BaseAgent]:
+def get_agent_class_from_agent_name(agent: str) -> type[BaseAgent]:
     if agent not in AGENT_MAPPING:
         raise ValueError(f"Agent '{agent}' is not recognized. Available agents: {list(AGENT_MAPPING.keys())}")
     return AGENT_MAPPING[agent]
diff --git a/src/lightman_ai/cli.py b/src/lightman_ai/cli.py
@@ -25,6 +25,7 @@ def entry_point() -> None:
     help=(f"Location of the config file containing the prompts. Defaults to `{DEFAULT_CONFIG_FILE}`."),
 )
 @click.option("--prompt", type=str, help=("Which prompt to use"))
+@click.option("--model", type=str, default=None, help=("Which model to use. Must be set in conjunction with --agent."))
 @click.option(
     "--score",
     type=int,
@@ -60,6 +61,7 @@ def run(
     agent: str,
     prompt: str,
     prompt_file: str,
+    model: str | None,
     score: int | None,
     config_file: str,
     config: str,
@@ -80,6 +82,7 @@ def run(
                 "agent": agent or config_from_file.agent,
                 "prompt": prompt or config_from_file.prompt,
                 "score_threshold": score or config_from_file.score_threshold,
+                "model": model or config_from_file.model,
             }
         )
 
@@ -94,6 +97,7 @@ def run(
         dry_run=dry_run,
         project_key=config_from_file.service_desk_project_key,
         request_id_type=config_from_file.service_desk_request_id_type,
+        model=final_config.model,
     )
     relevant_articles_metadata = [f"{article.title} ({article.link})" for article in relevant_articles]
     logger.warning("Found these articles: \n- %s", "\n- ".join(relevant_articles_metadata))
diff --git a/src/lightman_ai/core/config.py b/src/lightman_ai/core/config.py
@@ -44,6 +44,7 @@ class FileConfig(BaseModel):
     score_threshold: int | None = None
     service_desk_project_key: str | None = None
     service_desk_request_id_type: str | None = None
+    model: str | None = None
 
     model_config = ConfigDict(extra="forbid")
 
@@ -77,6 +78,7 @@ class FinalConfig(BaseModel):
     prompt: str
     agent: str
     score_threshold: PositiveInt
+    model: str | None = None
 
     @classmethod
     def init_from_dict(cls, data: dict[str, Any]) -> Self:
diff --git a/src/lightman_ai/main.py b/src/lightman_ai/main.py
@@ -2,7 +2,7 @@
 import logging
 
 from lightman_ai.ai.base.agent import BaseAgent
-from lightman_ai.ai.utils import get_agent_instance_from_agent_name
+from lightman_ai.ai.utils import get_agent_class_from_agent_name
 from lightman_ai.article.models import ArticlesList, SelectedArticle, SelectedArticlesList
 from lightman_ai.integrations.service_desk.integration import (
     ServiceDeskIntegration,
@@ -59,11 +59,12 @@ def lightman(
     project_key: str | None = None,
     request_id_type: str | None = None,
     dry_run: bool = False,
+    model: str | None = None,
 ) -> list[SelectedArticle]:
     articles: ArticlesList = _get_articles()
 
-    agent_class = get_agent_instance_from_agent_name(agent)
-    agent_instance = agent_class(prompt)
+    agent_class = get_agent_class_from_agent_name(agent)
+    agent_instance = agent_class(prompt, model, logger=logger)
     logger.info("Selected %s.", agent_instance)
 
     classified_articles = _classify_articles(
diff --git a/tests/ai/base/test_agent.py b/tests/ai/base/test_agent.py
@@ -1,23 +1,40 @@
-import logging
 from typing import override
-from unittest.mock import Mock
+from unittest.mock import Mock, patch
 
 from lightman_ai.ai.base.agent import BaseAgent
 from lightman_ai.article.models import ArticlesList, SelectedArticlesList
 
 
 class FakeAgent(BaseAgent):
-    def __init__(self, system_prompt: str, logger: logging.Logger | None = None) -> None:
-        self.agent = Mock()
+    _class = Mock()
+    _default_model_name = "default model name"
 
     @override
     def _run_prompt(self, prompt: str) -> SelectedArticlesList:
         return SelectedArticlesList(articles=[])
 
 
 class TestBaseAgent:
-    def test__get_prompt_result(self, test_prompt: str, thn_news: ArticlesList) -> None:
+    @patch("lightman_ai.ai.base.agent.Agent")
+    def test__get_prompt_result(self, m_agent: Mock, test_prompt: str, thn_news: ArticlesList) -> None:
         """Check that we receive an instance of `SelectedArticlesList` when running the method."""
         agent = FakeAgent(test_prompt)
-        result = agent.get_prompt_result(str(thn_news))
-        assert isinstance(result, SelectedArticlesList)
+
+        with patch("tests.ai.base.test_agent.FakeAgent._run_prompt") as m_run_prompt:
+            agent.get_prompt_result(str(thn_news))
+
+        assert m_run_prompt.call_count == 1
+        assert m_run_prompt.call_args[0][0] == str(thn_news)
+        assert m_agent.call_count == 1
+        assert m_agent.call_args[1]["system_prompt"] == test_prompt
+        assert agent._class.call_count == 1
+        assert agent._class.call_args[0][0] == FakeAgent._default_model_name
+
+    @patch("lightman_ai.ai.base.agent.Agent")
+    def test_agent_is_intantiated_with_model_when_set(self, m_agent: Mock, test_prompt: str) -> None:
+        agent = FakeAgent(test_prompt, model="my model")
+        agent.get_prompt_result("")
+
+        assert m_agent.call_count == 1
+        assert agent._class.call_count == 2
+        assert agent._class.call_args[0][0] == "my model"
diff --git a/tests/core/test_config.py b/tests/core/test_config.py
@@ -73,6 +73,7 @@ def test_file_is_loaded_from_location(self, tmp_path: pathlib.Path) -> None:
         assert config.agent == "openai"
         assert config.score_threshold == 8
         assert config.prompt == "eval-prompt"
+        assert config.model is None
 
     def test_service_desk_fields_accept_str(self) -> None:
         config = FileConfig(
@@ -128,6 +129,20 @@ def test_service_desk_fields_accept_none(self) -> None:
         assert config.service_desk_project_key is None
         assert config.service_desk_request_id_type is None
 
+    def test_model_parameter_can_be_set_in_config_file(self) -> None:
+        """Test that the model parameter can be set through the config file."""
+        content = """
+        [default]
+        agent = 'openai'
+        score_threshold = 8
+        prompt = 'eval-prompt'
+        model = 'gpt-4o-custom'
+        """
+        with patch_config_file(content=content):
+            config = FileConfig.get_config_from_file(config_section=DEFAULT_CONFIG_SECTION, path=DEFAULT_CONFIG_FILE)
+
+        assert config.model == "gpt-4o-custom"
+
 
 class TestFinalConfig:
     def test_init_error(self) -> None:
@@ -148,6 +163,23 @@ def test_score_must_be_positive_int(self) -> None:
             )
         assert "`score_threshold`: Input should be greater than 0" in exc.value.args[0]
 
+    def test_final_config_accepts_model_parameter(self) -> None:
+        """Test that FinalConfig now accepts a model parameter."""
+        config = FinalConfig.init_from_dict(
+            {"prompt": "test prompt", "agent": "openai", "score_threshold": 5, "model": "gpt-4o-custom"}
+        )
+
+        assert config.prompt == "test prompt"
+        assert config.agent == "openai"
+        assert config.score_threshold == 5
+        assert config.model == "gpt-4o-custom"
+
+    def test_final_config_model_parameter_is_optional(self) -> None:
+        """Test that the model parameter is optional and defaults to None."""
+        config = FinalConfig.init_from_dict({"prompt": "test prompt", "agent": "openai", "score_threshold": 5})
+
+        assert config.model is None
+
 
 class TestPromptConfig:
     def test_wrong_prompt(self) -> None:
diff --git a/tests/eval/test_cli.py b/tests/eval/test_cli.py
diff --git a/tests/test_cli.py b/tests/test_cli.py
