Merge pull request #24 from Deeptechia/develop

fix gemini markdown #23

geppetto/gemini_handler.py

@@ -1,3 +1,4 @@
+import re
 from urllib.request import urlopen
 import logging
 
@@ -19,8 +20,11 @@
 MSG_INPUT_FIELD = "content"
 
 def to_markdown(text):
-    text = text.replace('•', '  *')
-    return Markdown(textwrap.indent(text, '> ', predicate=lambda _: True))
+    formatted_text = text.replace("**", "*")
+    formatted_text = formatted_text.replace("__", "_")
+    formatted_text = formatted_text.replace("- ", "• ")
+    formatted_text = re.sub(r"\[(.*?)\]\((.*?)\)", r"<\2|\1>", formatted_text) 
+    return formatted_text
 
 class GeminiHandler(LLMHandler):
 
@@ -49,7 +53,7 @@ def llm_generate_content(self, user_prompt, status_callback=None, *status_callba
             user_prompt = [merged_prompt] + user_prompt[2:]
         response= self.client.generate_content(user_prompt)
         markdown_response = to_markdown(response.text)
-        return str(markdown_response.data)
+        return markdown_response
 
     def get_prompt_from_thread(self, thread: List[Dict], assistant_tag: str, user_tag: str):
         prompt = super().get_prompt_from_thread(thread, assistant_tag, user_tag)

tests/test_gemini.py

@@ -41,11 +41,11 @@ def test_llm_generate_content(self, mock_to_markdown):
         mock_response = Mock()
         mock_response.text = "Mocked Gemini response"
         self.gemini_handler.client.generate_content.return_value = mock_response
-        mock_to_markdown.return_value.data = "Mocked Markdown data"
+        mock_to_markdown.return_value = Mock(data="Mocked Markdown data")
 
         response = self.gemini_handler.llm_generate_content(user_prompt)
 
-        self.assertEqual(response, "Mocked Markdown data")
+        self.assertEqual(response.data, "Mocked Markdown data")
         mock_to_markdown.assert_called_once_with("Mocked Gemini response")
 
     def test_get_prompt_from_thread(self):