OLS-2673: Add query mode field to support different response styles

docs/openapi.json

@@ -1482,6 +1482,10 @@
                             }
                         ],
                         "title": "Mcp Headers"
+                    },
+                    "mode": {
+                        "$ref": "#/components/schemas/QueryMode",
+                        "default": "ask"
                     }
                 },
                 "additionalProperties": false,
@@ -1490,7 +1494,7 @@
                     "query"
                 ],
                 "title": "LLMRequest",
-                "description": "Model representing a request for the LLM (Language Model) send into OLS service.\n\nAttributes:\n    query: The query string.\n    conversation_id: The optional conversation ID (UUID).\n    provider: The optional provider.\n    model: The optional model.\n    attachments: The optional attachments.\n    media_type: The optional parameter for streaming response.\n    mcp_headers: Optional JSON object mapping MCP server names to header lists.\n\nExample:\n    ```python\n    llm_request = LLMRequest(query=\"Tell me about Kubernetes\")\n    ```",
+                "description": "Model representing a request for the LLM (Language Model) send into OLS service.\n\nAttributes:\n    query: The query string.\n    conversation_id: The optional conversation ID (UUID).\n    provider: The optional provider.\n    model: The optional model.\n    attachments: The optional attachments.\n    media_type: The optional parameter for streaming response.\n    mcp_headers: Optional JSON object mapping MCP server names to header lists.\n    mode: The query mode controlling which system prompt is used.\n\nExample:\n    ```python\n    llm_request = LLMRequest(query=\"Tell me about Kubernetes\")\n    ```",
                 "examples": [
                     {
                         "attachments": [
@@ -1513,6 +1517,7 @@
                         "conversation_id": "123e4567-e89b-12d3-a456-426614174000",
                         "mcp_headers": "{\"github-mcp\": {\"Authorization\": \"Bearer ghp_xxxxx\"}}",
                         "media_type": "text/plain",
+                        "mode": "ask",
                         "model": "model-name",
                         "provider": "openai",
                         "query": "write a deployment yaml for the mongodb image",
@@ -1910,6 +1915,15 @@
                     }
                 ]
             },
+            "QueryMode": {
+                "type": "string",
+                "enum": [
+                    "ask",
+                    "troubleshooting"
+                ],
+                "title": "QueryMode",
+                "description": "Query modes that control which system prompt is used."
+            },
             "ReadinessResponse": {
                 "properties": {
                     "ready": {

ols/app/endpoints/ols.py

@@ -378,6 +378,7 @@ def process_request(auth: Any, llm_request: LLMRequest) -> ProcessedRequest:
         timestamps=timestamps,
         skip_user_id_check=skip_user_id_check,
         user_token=user_token,
+        mode=llm_request.mode,
     )
 
 
@@ -560,6 +561,7 @@ def generate_response(
             provider=llm_request.provider,
             model=llm_request.model,
             system_prompt=llm_request.system_prompt,
+            mode=llm_request.mode,
             user_token=user_token,
             client_headers=client_headers,
         )
@@ -806,6 +808,7 @@ def store_transcript(
             "model": llm_request.model or config.ols_config.default_model,
             "user_id": user_id,
             "conversation_id": conversation_id,
+            "mode": llm_request.mode,
             "timestamp": datetime.now(pytz.UTC).isoformat(),
         },
         "redacted_query": redacted_query,

ols/app/models/models.py

@@ -10,7 +10,7 @@
 from pydantic import BaseModel, Field, field_validator, model_validator
 from pydantic.dataclasses import dataclass
 
-from ols.constants import MEDIA_TYPE_JSON, MEDIA_TYPE_TEXT
+from ols.constants import MEDIA_TYPE_JSON, MEDIA_TYPE_TEXT, QueryMode
 from ols.src.prompts.prompts import QUERY_SYSTEM_INSTRUCTION
 from ols.utils import suid
 
@@ -73,6 +73,7 @@ class LLMRequest(BaseModel):
         attachments: The optional attachments.
         media_type: The optional parameter for streaming response.
         mcp_headers: Optional JSON object mapping MCP server names to header lists.
+        mode: The query mode controlling which system prompt is used.
 
     Example:
         ```python
@@ -88,6 +89,7 @@ class LLMRequest(BaseModel):
     attachments: Optional[list[Attachment]] = None
     media_type: Optional[str] = MEDIA_TYPE_TEXT
     mcp_headers: Optional[dict[str, dict[str, str]]] = None
+    mode: QueryMode = QueryMode.ASK
 
     # provides examples for /docs endpoint
     model_config = {
@@ -119,6 +121,7 @@ class LLMRequest(BaseModel):
                     ],
                     "media_type": "text/plain",
                     "mcp_headers": '{"github-mcp": {"Authorization": "Bearer ghp_xxxxx"}}',
+                    "mode": "ask",
                 }
             ]
         },
@@ -970,6 +973,7 @@ class ProcessedRequest(BaseModel):
         timestamps: Timestamps for all operations.
         skip_user_id_check: Flag to skip user ID checking in handler.
         user_token: User token (if provided).
+        mode: The query mode controlling which system prompt is used.
     """
 
     user_id: str
@@ -980,6 +984,7 @@ class ProcessedRequest(BaseModel):
     timestamps: dict[str, float]
     skip_user_id_check: bool
     user_token: str
+    mode: QueryMode
 
 
 @dataclass

ols/constants.py

@@ -34,6 +34,13 @@ class ModelFamily(StrEnum):
     GRANITE = "granite"
 
 
+class QueryMode(StrEnum):
+    """Query modes that control which system prompt is used."""
+
+    ASK = "ask"
+    TROUBLESHOOTING = "troubleshooting"
+
+
 class GenericLLMParameters:
     """Generic LLM parameters that can be mapped into LLM provider-specific parameters."""
 

ols/src/prompts/prompts.py

@@ -84,6 +84,8 @@
 * Terseness must not omit critical info.
 """
 
+TROUBLESHOOTING_SYSTEM_INSTRUCTION = "NOT IMPLEMENTED"
+
 USE_CONTEXT_INSTRUCTION = """
 Use the retrieved document to answer the question.
 """

ols/src/query_helpers/query_helper.py

@@ -7,11 +7,20 @@
 from langchain_core.language_models.llms import LLM
 
 from ols import config
+from ols.constants import QueryMode
 from ols.src.llms.llm_loader import load_llm
-from ols.src.prompts.prompts import QUERY_SYSTEM_INSTRUCTION
+from ols.src.prompts.prompts import (
+    QUERY_SYSTEM_INSTRUCTION,
+    TROUBLESHOOTING_SYSTEM_INSTRUCTION,
+)
 
 logger = logging.getLogger(__name__)
 
+_DEFAULT_PROMPT_BY_MODE = {
+    QueryMode.ASK: QUERY_SYSTEM_INSTRUCTION,
+    QueryMode.TROUBLESHOOTING: TROUBLESHOOTING_SYSTEM_INSTRUCTION,
+}
+
 
 class QueryHelper:
     """Base class for query helpers."""
@@ -23,6 +32,7 @@ def __init__(
         generic_llm_params: Optional[dict] = None,
         llm_loader: Optional[Callable[[str, str, dict], LLM]] = None,
         system_prompt: Optional[str] = None,
+        mode: QueryMode = QueryMode.ASK,
     ) -> None:
         """Initialize query helper."""
         # NOTE: As signature of this method is evaluated before the config,
@@ -36,6 +46,6 @@ def __init__(
         self._system_prompt = (
             (config.dev_config.enable_system_prompt_override and system_prompt)
             or config.ols_config.system_prompt
-            or QUERY_SYSTEM_INSTRUCTION
+            or _DEFAULT_PROMPT_BY_MODE[mode]
         )
         logger.debug("System prompt: %s", self._system_prompt)

tests/config/config_for_integration_tests.yaml

@@ -11,7 +11,7 @@ llm_providers:
       - name: m1
         url: "https://murl1"
         credentials_path: tests/config/secret/apitoken
-        context_window_size: 450
+        context_window_size: 1000
         parameters:
           max_tokens_for_response: 100
       - name: m2