moves prompt_modules to utils

Author: Tulio Coppola

cli/decompose/decompose.py

@@ -5,6 +5,8 @@
 
 import typer
 
+from .pipeline import DecompBackend
+
 this_file_dir = Path(__file__).resolve().parent
 
 
@@ -20,15 +22,69 @@ def run(
         typer.FileText | None,
         typer.Option(help="Path to a raw text file containing a task prompt."),
     ] = None,
+    model_id: Annotated[
+        str,
+        typer.Option(
+            help=(
+                "Model name/id to be used to run the decomposition pipeline."
+                + ' Defaults to "mistral-small3.2:latest", which is valid for the "ollama" backend.'
+                + " If you have a vLLM instance serving a model from HF with vLLM's OpenAI"
+                + " compatible endpoint, then this option should be set to the model's HF name/id,"
+                + ' e.g. "mistralai/Mistral-Small-3.2-24B-Instruct-2506" and the "--backend" option'
+                + ' should be set to "openai".'
+            )
+        ),
+    ] = "mistral-small3.2:latest",
+    backend: Annotated[
+        DecompBackend,
+        typer.Option(
+            help=(
+                'Backend to be used for inference. Defaults to "ollama".'
+                + ' Options are: "ollama" and "openai".'
+                + ' The "ollama" backend runs a local inference server.'
+                + ' The "openai" backend will send inference requests to any'
+                + " endpoint that's OpenAI compatible."
+            ),
+            case_sensitive=False,
+        ),
+    ] = DecompBackend.ollama,
+    backend_req_timeout: Annotated[
+        int,
+        typer.Option(
+            help='Time (in seconds) for timeout to be passed on the model inference requests. Defaults to "3600"'
+        ),
+    ] = 3600,
+    backend_endpoint: Annotated[
+        str | None,
+        typer.Option(
+            help=(
+                'The "endpoint URL", sometimes called "base URL",'
+                + ' to reach the model when using the "openai" backend.'
+                + ' This option is required if using "--backend openai".'
+            )
+        ),
+    ] = None,
+    backend_api_key: Annotated[
+        str | None,
+        typer.Option(
+            help=(
+                'The API key for the configured "--backend-endpoint".'
+                + ' If using "--backend openai" this option must be set,'
+                + " even if you are running locally (an OpenAI compatible server), you"
+                + ' must set this option, it can be set to "EMPTY" if your local'
+                + " server doesn't need it."
+            )
+        ),
+    ] = None,
     input_var: Annotated[
         list[str] | None,
         typer.Option(
             help=(
-                "If your task prompt needs user input data, you must pass"
-                + " a descriptive variable name using this option,"
-                + " so the name can be included when generating the prompt templates."
+                "If your task needs user input data, you must pass"
+                + " a descriptive variable name using this option, this way"
+                + " the variable names can be templated into the generated prompts."
                 + " You can pass this option multiple times, one for each input variable name."
-                + " These names must be all uppercase, alphanumeric with words separated by underscores."
+                + " These names must be all uppercase, alphanumeric, with words separated by underscores."
             )
         ),
     ] = None,
@@ -63,7 +119,13 @@ def run(
 
         if prompt_file:
             decomp_data = pipeline.decompose(
-                task_prompt=prompt_file.read(), user_input_variable=input_var
+                task_prompt=prompt_file.read(),
+                user_input_variable=input_var,
+                model_id=model_id,
+                backend=backend,
+                backend_req_timeout=backend_req_timeout,
+                backend_endpoint=backend_endpoint,
+                backend_api_key=backend_api_key,
             )
         else:
             task_prompt: str = typer.prompt(
@@ -76,7 +138,13 @@ def run(
             )
             task_prompt = task_prompt.replace("\\n", "\n")
             decomp_data = pipeline.decompose(
-                task_prompt=task_prompt, user_input_variable=None
+                task_prompt=task_prompt,
+                user_input_variable=None,
+                model_id=model_id,
+                backend=backend,
+                backend_req_timeout=backend_req_timeout,
+                backend_endpoint=backend_endpoint,
+                backend_api_key=backend_api_key,
             )
 
         with open(out_dir / f"{out_name}.json", "w") as f:

cli/decompose/pipeline.py

@@ -1,19 +1,23 @@
+from enum import Enum
 from typing import TypedDict
 
 from typing_extensions import NotRequired
 
 from mellea import MelleaSession
-from mellea.backends import ModelOption
 from mellea.backends.ollama import OllamaModelBackend
-from mellea.prompt_modules import (
+from mellea.backends.openai import OpenAIBackend
+from mellea.backends.types import ModelOption
+from mellea.helpers.prompt_modules import (
     constraint_extractor,
     subtask_constraint_assign,
     subtask_list,
     subtask_prompt_generator,
 )
-from mellea.prompt_modules.subtask_constraint_assign import SubtaskPromptConstraintsItem
-from mellea.prompt_modules.subtask_list import SubtaskItem
-from mellea.prompt_modules.subtask_prompt_generator import SubtaskPromptItem
+from mellea.helpers.prompt_modules.subtask_constraint_assign import (
+    SubtaskPromptConstraintsItem,
+)
+from mellea.helpers.prompt_modules.subtask_list import SubtaskItem
+from mellea.helpers.prompt_modules.subtask_prompt_generator import SubtaskPromptItem
 
 
 class DecompSubtasksResult(TypedDict):
@@ -32,37 +36,84 @@ class DecompPipelineResult(TypedDict):
     final_response: NotRequired[str]
 
 
+class DecompBackend(str, Enum):
+    ollama = "ollama"
+    openai = "openai"
+    rits = "rits"
+
+
 def decompose(
-    task_prompt: str, user_input_variable: list[str] | None = None
+    task_prompt: str,
+    user_input_variable: list[str] | None = None,
+    model_id: str = "mistral-small3.2:latest",
+    backend: DecompBackend = DecompBackend.ollama,
+    backend_req_timeout: int = 3600,
+    backend_endpoint: str | None = None,
+    backend_api_key: str | None = None,
 ) -> DecompPipelineResult:
     if user_input_variable is None:
         user_input_variable = []
 
-    m_ollama_session = MelleaSession(
-        OllamaModelBackend(
-            model_id="mistral-small3.2:24b",
-            model_options={ModelOption.CONTEXT_WINDOW: 32768},
-        )
-    )
+    match backend:
+        case DecompBackend.ollama:
+            m_session = MelleaSession(
+                OllamaModelBackend(
+                    model_id=model_id,
+                    model_options={
+                        ModelOption.CONTEXT_WINDOW: 32768,
+                        "timeout": backend_req_timeout,
+                    },
+                )
+            )
+        case DecompBackend.openai:
+            assert backend_endpoint is not None, (
+                'Required to provide "backend_endpoint" for this configuration'
+            )
+            assert backend_api_key is not None, (
+                'Required to provide "backend_api_key" for this configuration'
+            )
+            m_session = MelleaSession(
+                OpenAIBackend(
+                    model_id=model_id,
+                    base_url=backend_endpoint,
+                    api_key=backend_api_key,
+                    model_options={"timeout": backend_req_timeout},
+                )
+            )
+        case DecompBackend.rits:
+            assert backend_endpoint is not None, (
+                'Required to provide "backend_endpoint" for this configuration'
+            )
+            assert backend_api_key is not None, (
+                'Required to provide "backend_api_key" for this configuration'
+            )
 
-    subtasks: list[SubtaskItem] = subtask_list.generate(
-        m_ollama_session, task_prompt
-    ).parse()
+            from mellea_ibm.rits import RITSBackend, RITSModelIdentifier  # type: ignore
+
+            m_session = MelleaSession(
+                RITSBackend(
+                    RITSModelIdentifier(endpoint=backend_endpoint, model_name=model_id),
+                    api_key=backend_api_key,
+                    model_options={"timeout": backend_req_timeout},
+                )
+            )
+
+    subtasks: list[SubtaskItem] = subtask_list.generate(m_session, task_prompt).parse()
 
     task_prompt_constraints: list[str] = constraint_extractor.generate(
-        m_ollama_session, task_prompt
+        m_session, task_prompt
     ).parse()
 
     subtask_prompts: list[SubtaskPromptItem] = subtask_prompt_generator.generate(
-        m_ollama_session,
+        m_session,
         task_prompt,
         user_input_var_names=user_input_variable,
         subtasks_and_tags=subtasks,
     ).parse()
 
     subtask_prompts_with_constraints: list[SubtaskPromptConstraintsItem] = (
         subtask_constraint_assign.generate(
-            m_ollama_session,
+            m_session,
             subtasks_tags_and_prompts=subtask_prompts,
             constraint_list=task_prompt_constraints,
         ).parse()

mellea/prompt_modules/__init__.py renamed to mellea/helpers/prompt_modules/__init__.py

@@ -4,7 +4,10 @@
 
 from mellea import MelleaSession
 from mellea.backends.types import ModelOption
-from mellea.prompt_modules._prompt_modules import PromptModule, PromptModuleString
+from mellea.helpers.prompt_modules._prompt_modules import (
+    PromptModule,
+    PromptModuleString,
+)
 from mellea.stdlib.instruction import Instruction
 
 from ._exceptions import BackendGenerationError, TagExtractionError