Fix Selected Response Bug

CHANGELOG.md

@@ -8,6 +8,12 @@ and this project adheres to [Semantic Versioning](http://semver.org/spec/v2.0.0.
 ### Changed
 - `Python 3.15` added to `test.yml`
 - Build system modified
+- Test system modified
+- Selected response bug fixed
+- `Prompt` class `remove_response` method modified
+- `Prompt` class `clear_responses` method modified
+- `Prompt` class `select_response` method modified
+- `Prompt` class `update_responses` method modified
 ## [1.4] - 2026-10-05
 ### Added
 - Z.AI models

memor/params.py

@@ -18,6 +18,7 @@
 INVALID_DATETIME_MESSAGE = "Invalid value. `{parameter_name}` must be a datetime object that includes timezone information."
 INVALID_TEMPLATE_MESSAGE = "Invalid template. It must be an instance of `PromptTemplate` or `PresetPromptTemplate`."
 INVALID_RESPONSE_MESSAGE = "Invalid response. It must be an instance of `Response`."
+INVALID_RESPONSE_IDENTIFIER_MESSAGE = "Invalid response identifier. It must be a non-negative integer index or a valid response ID."
 INVALID_MESSAGE = "Invalid message. It must be an instance of `Prompt` or `Response`."
 INVALID_MESSAGE_STATUS_LEN_MESSAGE = "Invalid message status length. It must be equal to the number of messages."
 INVALID_CUSTOM_MAP_MESSAGE = "Invalid custom map: it must be a dictionary with keys and values that can be converted to strings."

memor/prompt.py

@@ -3,14 +3,15 @@
 from typing import List, Dict, Union, Tuple, Any, Optional
 import datetime
 import json
+import uuid
 import warnings
 from .message import Message
 from .params import MEMOR_VERSION
 from .params import DATE_TIME_FORMAT
 from .params import RenderFormat, DATA_SAVE_SUCCESS_MESSAGE
 from .params import Role
 from .params import INVALID_PROMPT_STRUCTURE_MESSAGE, INVALID_TEMPLATE_MESSAGE
-from .params import INVALID_RESPONSE_MESSAGE
+from .params import INVALID_RESPONSE_MESSAGE, INVALID_RESPONSE_IDENTIFIER_MESSAGE
 from .params import PROMPT_RENDER_ERROR_MESSAGE
 from .params import INVALID_RENDER_FORMAT_MESSAGE
 from .params import AI_STUDIO_SYSTEM_WARNING
@@ -60,7 +61,7 @@ def __init__(
         self._role = Role.USER
         self._template = PresetPromptTemplate.DEFAULT.value
         self._responses = []
-        self._selected_response_index = 0
+        self._selected_response_id = None
         if file_path is not None:
             self.load(file_path)
         else:
@@ -112,30 +113,50 @@ def add_response(self, response: Response, index: Optional[int] = None) -> None:
 
     def remove_response(self, index: int) -> None:
         """
-        Remove a response from the prompt object.
+        Remove a response from the prompt.
 
         :param index: index
         """
-        self._responses.pop(index)
+        response = self._responses.pop(index)
+        if response.id == self._selected_response_id:
+            self._selected_response_id = None
         self._mark_modified()
 
     def clear_responses(self) -> None:
         """Remove all responses from the prompt."""
         self._responses.clear()
-        self._selected_response_index = 0
+        self._selected_response_id = None
         self._mark_modified()
 
-    def select_response(self, index: int) -> Optional[Response]:
+    def select_response(self, identifier: Optional[Union[int, str]]) -> Optional[Response]:
         """
-        Select a response as selected response.
+        Select a response by index or UUIDv4 ID.
 
-        :param index: index
+        :param identifier: non-negative index or response UUIDv4
         """
-        _validate_non_negative_int(index, "index")
-        self._selected_response_index = index
+        if identifier is None:
+            response = None
+        elif isinstance(identifier, int):
+            _validate_non_negative_int(identifier, "identifier")
+            if identifier >= len(self._responses):
+                response = None
+            else:
+                response = self._responses[identifier]
+        elif isinstance(identifier, str):
+            try:
+                _ = uuid.UUID(identifier, version=4)
+            except Exception:
+                raise MemorValidationError(INVALID_RESPONSE_IDENTIFIER_MESSAGE)
+            response = next((item for item in self._responses if item.id == identifier), None)
+        else:
+            raise MemorValidationError(INVALID_RESPONSE_IDENTIFIER_MESSAGE)
+
+        if response is None:
+            self._selected_response_id = None
+        else:
+            self._selected_response_id = response.id
         self._mark_modified()
-        if index < len(self._responses):
-            return self._responses[index]
+        return response
 
     def update_responses(self, responses: List[Response]) -> None:
         """
@@ -144,7 +165,11 @@ def update_responses(self, responses: List[Response]) -> None:
         :param responses: responses
         """
         _validate_list_of(responses, "responses", Response, "`Response`")
-        self._responses = responses
+        self._responses = list(responses)
+
+        if not any(response.id == self._selected_response_id for response in self._responses):
+            self._selected_response_id = None
+
         self._mark_modified()
 
     def update_template(self, template: Union[PresetPromptTemplate, PromptTemplate]) -> None:
@@ -206,17 +231,24 @@ def _validate_extract_json(json_object: Union[str, Dict[str, Any]]) -> Dict[str,
             result["memor_version"] = loaded_obj["memor_version"]
             result["date_created"] = datetime.datetime.strptime(loaded_obj["date_created"], DATE_TIME_FORMAT)
             result["date_modified"] = datetime.datetime.strptime(loaded_obj["date_modified"], DATE_TIME_FORMAT)
-            result["selected_response_index"] = loaded_obj["selected_response_index"]
+            result["selected_response_id"] = loaded_obj.get("selected_response_id")
+            result["selected_response_index"] = loaded_obj.get("selected_response_index")
         except Exception:
             raise MemorValidationError(INVALID_PROMPT_STRUCTURE_MESSAGE)
         _validate_string(result["message"], "message")
 
         if result["tokens"] is not None:
             _validate_non_negative_int(result["tokens"], "tokens")
+        if result["selected_response_id"] is not None:
+            try:
+                _ = uuid.UUID(result["selected_response_id"], version=4)
+            except Exception:
+                raise MemorValidationError(INVALID_RESPONSE_IDENTIFIER_MESSAGE)
         _validate_message_id(result["id"])
         _validate_warnings(result["warnings"])
         _validate_string(result["memor_version"], "memor_version")
-        _validate_non_negative_int(result["selected_response_index"], "selected_response_index")
+        if result["selected_response_index"] is not None:
+            _validate_non_negative_int(result["selected_response_index"], "selected_response_index")
         return result
 
     def from_json(self, json_object: Union[str, Dict[str, Any]]) -> None:
@@ -236,7 +268,8 @@ def from_json(self, json_object: Union[str, Dict[str, Any]]) -> None:
         self._memor_version = data["memor_version"]
         self._date_created = data["date_created"]
         self._date_modified = data["date_modified"]
-        self.select_response(data["selected_response_index"])
+        selected_response_identifier = data["selected_response_id"] or data["selected_response_index"]
+        self.select_response(selected_response_identifier)
 
     def to_json(self, save_template: bool = True) -> Dict[str, Any]:
         """
@@ -265,7 +298,7 @@ def to_dict(self, save_template: bool = True) -> Dict[str, Any]:
             "message": self._message,
             "warnings": self._warnings,
             "responses": self._responses.copy(),
-            "selected_response_index": self._selected_response_index,
+            "selected_response_id": self._selected_response_id,
             "tokens": self._tokens,
             "role": self._role,
             "id": self._id,
@@ -291,8 +324,9 @@ def template(self) -> PromptTemplate:
     @property
     def selected_response(self) -> Optional[Response]:
         """Get the prompt selected response."""
-        if 0 <= self._selected_response_index < len(self._responses):
-            return self._responses[self._selected_response_index]
+        for response in self._responses:
+            if response.id == self._selected_response_id:
+                return response
         return None
 
     def render(self, render_format: RenderFormat = RenderFormat.DEFAULT,
@@ -311,6 +345,9 @@ def render(self, render_format: RenderFormat = RenderFormat.DEFAULT,
             format_kwargs = {"prompt": self.to_json(save_template=False)}
             if isinstance(self.selected_response, Response):
                 format_kwargs.update({"response": self.selected_response.to_json()})
+            else:
+                if len(self.responses) > 0 and isinstance(self.responses[0], Response):
+                    format_kwargs.update({"response": self.responses[0].to_json()})
             responses_dicts = []
             for _, response in enumerate(self._responses):
                 responses_dicts.append(response.to_json())

tests/test_prompt.py

@@ -241,7 +241,7 @@ def test_clear_responses():
     response0 = Response(message="I am fine.")
     response1 = Response(message="Great!")
     prompt = Prompt(message=message, responses=[response0, response1])
-    _ = prompt.select_response(index=0)
+    _ = prompt.select_response(0)
     assert len(prompt.responses) == 2
     assert prompt.selected_response == response0
     prompt.clear_responses()
@@ -255,7 +255,7 @@ def test_select_response():
     prompt = Prompt(message=message, responses=[response0])
     response1 = Response(message="Great!")
     prompt.add_response(response1)
-    selected_response = prompt.select_response(index=1)
+    selected_response = prompt.select_response(1)
     assert prompt.selected_response == response1
     assert selected_response == response1
 
@@ -265,10 +265,131 @@ def test_select_response2():
     response0 = Response("I am fine.")
     response1 = Response("Great!")
     prompt.update_responses([response0, response1])
-    selected_response = prompt.select_response(20)
-    assert prompt._selected_response_index == 20
+    assert prompt.select_response(20) is None
     assert prompt.selected_response is None
-    assert selected_response is None
+
+
+def test_select_response3():
+    message = "Hello, how are you?"
+    response0 = Response(message="I am fine.")
+    response1 = Response(message="Great!")
+    prompt = Prompt(message=message, responses=[response0, response1])
+    selected_response = prompt.select_response(response1.id)
+    assert prompt.selected_response == response1
+    assert selected_response == response1
+
+
+def test_select_response4():
+    response0 = Response(message="I am fine.")
+    response1 = Response(message="Great!")
+    response2 = Response(message="Thanks!")
+    prompt = Prompt(responses=[response0, response1, response2])
+    prompt.select_response(response1.id)
+    prompt.update_responses([response2, response0, response1])
+    assert prompt.selected_response == response1
+
+
+def test_select_response5():
+    response0 = Response(message="I am fine.")
+    response1 = Response(message="Great!")
+    response2 = Response(message="Thanks!")
+    prompt = Prompt(responses=[response0, response1, response2])
+    prompt.select_response(response1.id)
+    prompt.remove_response(0)
+    assert prompt.selected_response == response1
+
+
+def test_select_response6():
+    response0 = Response(message="I am fine.")
+    response1 = Response(message="Great!")
+    prompt = Prompt(responses=[response0, response1])
+    prompt.select_response(response1.id)
+    prompt.remove_response(1)
+    assert prompt.selected_response is None
+
+
+def test_select_response7():
+    response0 = Response(message="I am fine.")
+    response1 = Response(message="Great!")
+    prompt = Prompt(responses=[response0, response1])
+    prompt.select_response(response1.id)
+    prompt_json = prompt.to_json()
+    assert prompt_json["selected_response_id"] == response1.id
+    assert "selected_response_index" not in prompt_json
+
+
+def test_select_response8():
+    response0 = Response(message="I am fine.")
+    response1 = Response(message="Great!")
+    prompt = Prompt(responses=[response0, response1])
+    prompt.select_response(1)
+    prompt_json = prompt.to_json()
+    prompt_json.pop("selected_response_id")
+    prompt_json["selected_response_index"] = 1
+    loaded_prompt = Prompt()
+    loaded_prompt.from_json(prompt_json)
+    assert loaded_prompt.selected_response == response1
+
+
+@pytest.mark.parametrize(
+    "identifier",
+    [
+        -1,
+        True,
+        False,
+        1.5,
+        [],
+        {},
+        "invalid-uuid",
+    ],
+)
+def test_select_response9(identifier):
+    prompt = Prompt(responses=[Response(message="I am fine.")])
+    with pytest.raises(MemorValidationError):
+        prompt.select_response(identifier)
+
+
+def test_select_response10():
+    prompt = Prompt(responses=[Response(message="I am fine.")])
+    unknown_id = str(uuid.uuid4())
+    assert prompt.select_response(unknown_id) is None
+    assert prompt.selected_response is None
+
+
+def test_select_response11():
+    prompt = Prompt()
+    assert prompt.select_response(0) is None
+    assert prompt.selected_response is None
+    assert prompt.select_response(str(uuid.uuid4())) is None
+    assert prompt.selected_response is None
+
+
+def test_select_response12():
+    response0 = Response(message="I am fine.")
+    response1 = Response(message="Great!")
+    prompt = Prompt(responses=[response0, response1])
+    assert prompt.select_response(None) is None
+    assert prompt.selected_response is None
+
+
+def test_select_response13():
+    response0 = Response(message="I am fine.")
+    response1 = Response(message="Great!")
+    prompt = Prompt(responses=[response0, response1])
+    prompt.select_response(response1.id)
+    prompt_json = prompt.to_json()
+    loaded_prompt = Prompt()
+    loaded_prompt.from_json(prompt_json)
+    assert loaded_prompt.selected_response == response1
+
+
+def test_select_response14():
+    prompt = Prompt(message="Hello, how are you?")
+    prompt_json = prompt.to_json()
+    prompt_json["selected_response_id"] = "invalid_uuid"
+    loaded_prompt = Prompt()
+    with pytest.raises(MemorValidationError, match=r"Invalid response identifier. It must be a non-negative integer index or a valid response ID."):
+        loaded_prompt.from_json(prompt_json)
 
 
 def test_template1():
@@ -839,6 +960,21 @@ def test_render10():
     assert prompt.render(RenderFormat.AI_STUDIO) == {'role': 'model', 'parts': [{'text': 'Hi, How are you?'}]}
 
 
+def test_render11():
+    message = "Hello, how are you?"
+    response1 = Response(message="I am fine.", model=LLMModel.GPT_4, temperature=0.5, role=Role.USER, score=0.8)
+    response2 = Response(message="Thanks!", model=LLMModel.GPT_4, temperature=0.5, role=Role.USER, score=0.8)
+    prompt = Prompt(
+        message=message,
+        responses=[
+            response1,
+            response2],
+        role=Role.USER,
+        template=PresetPromptTemplate.BASIC.PROMPT_RESPONSE_STANDARD)
+    prompt.select_response(1)
+    assert prompt.render() == 'Prompt: Hello, how are you?\nResponse: Thanks!'
+
+
 def test_size_warning1():
     message = "How are you?"
     response1 = Response(message="I am fine.", model=LLMModel.GPT_4, temperature=0.5, role=Role.USER, score=0.8)