symbolicai 0.21.0__py3-none-any.whl → 1.1.0__py3-none-any.whl

symai/extended/api_builder.py

@@ -1,9 +1,9 @@
 from .. import core
-from ..pre_processors import PreProcessor
+from ..components import Execute
 from ..post_processors import CodeExtractPostProcessor
+from ..pre_processors import PreProcessor
 from ..symbol import Expression, Symbol
-from ..components import Execute
-
+from ..utils import UserMessage
 
 API_BUILDER_DESCRIPTION = """[Description]
 You are an API coding tool for Python that creates API calls to any web URL based on user requests.
@@ -64,7 +64,7 @@ res = run(value) # [MANAGED] must contain this line, do not change
 
 class APIBuilderPreProcessor(PreProcessor):
     def __call__(self, argument):
-        return '$> {} =>'.format(str(argument.args[0]))
+        return f"$> {argument.args[0]!s} =>"
 
 
 class APIBuilder(Expression):
@@ -77,9 +77,12 @@ class APIBuilder(Expression):
         self.sym_return_type = APIBuilder
 
     def forward(self, sym: Symbol, **kwargs) -> Symbol:
-        @core.zero_shot(prompt="Build the API call code:\n",
-                   pre_processors=[APIBuilderPreProcessor()],
-                   post_processors=[CodeExtractPostProcessor()], **kwargs)
+        @core.zero_shot(
+            prompt="Build the API call code:\n",
+            pre_processors=[APIBuilderPreProcessor()],
+            post_processors=[CodeExtractPostProcessor()],
+            **kwargs,
+        )
         def _func(_, text) -> str:
             pass
 
@@ -96,18 +99,20 @@ class StackTraceRetryExecutor(Expression):
     def forward(self, code: Symbol, request: Symbol, **kwargs) -> Symbol:
         code = str(code)
         # Set value that gets passed on to the 'run' function in the generated code
-        value = request.value # do not remove this line
+        value = request.value  # do not remove this line
         # Create the 'run' function
         self._runnable = self.executor(code, locals=locals().copy(), globals=globals().copy())
-        result = self._runnable['locals']['run'](value)
+        result = self._runnable["locals"]["run"](value)
         retry = 0
         # Retry if there is a 'Traceback' in the result
-        while 'Traceback' in result and retry <= self.max_retries:
-            self._runnable = self.executor(code, payload=result, locals=locals().copy(), globals=globals().copy(), **kwargs)
-            result = self._runnable['locals']['run'](value)
+        while "Traceback" in result and retry <= self.max_retries:
+            self._runnable = self.executor(
+                code, payload=result, locals=locals().copy(), globals=globals().copy(), **kwargs
+            )
+            result = self._runnable["locals"]["run"](value)
             retry += 1
-        if 'locals_res' in self._runnable:
-            result = self._runnable['locals_res']
+        if "locals_res" in self._runnable:
+            result = self._runnable["locals_res"]
         return result
 
 
@@ -126,14 +131,17 @@ class APIExecutor(Expression):
     def _runnable(self):
         return self.executor._runnable
 
-    def forward(self, request: Symbol, **kwargs) -> Symbol:
+    def forward(self, request: Symbol, **_kwargs) -> Symbol:
         self._request = self._to_symbol(request)
-        if self._verbose: print('[REQUEST]', self._request)
+        if self._verbose:
+            UserMessage(f"[REQUEST] {self._request}")
         # Generate the code to implement the API call
-        self._code    = self.builder(self._request)
-        if self._verbose: print('[GENERATED_CODE]', self._code)
+        self._code = self.builder(self._request)
+        if self._verbose:
+            UserMessage(f"[GENERATED_CODE] {self._code}")
         # Execute the code to define the 'run' function
-        self._result  = self.executor(self._code, request=self._request)
-        if self._verbose: print('[RESULT]:', self._result)
-        self._value   = self._result
+        self._result = self.executor(self._code, request=self._request)
+        if self._verbose:
+            UserMessage(f"[RESULT]: {self._result}")
+        self._value = self._result
         return self

symai/extended/arxiv_pdf_parser.py

@@ -1,17 +1,20 @@
-import os
 import re
 import shutil
-import requests
-
 from concurrent.futures import ThreadPoolExecutor, as_completed
+from pathlib import Path
+
+import requests
 
+from ..backend.settings import HOME_PATH
 from ..symbol import Expression, Symbol
+from ..utils import UserMessage
 from .file_merger import FileMerger
-from ..backend.settings import HOME_PATH
 
 
 class ArxivPdfParser(Expression):
-    def __init__(self, url_pattern: str = r'https://arxiv.org/(?:pdf|abs)/(\d+.\d+)(?:\.pdf)?', **kwargs):
+    def __init__(
+        self, url_pattern: str = r"https://arxiv.org/(?:pdf|abs)/(\d+.\d+)(?:\.pdf)?", **kwargs
+    ):
         super().__init__(**kwargs)
         self.url_pattern = url_pattern
         self.merger = FileMerger()
@@ -21,28 +24,34 @@ class ArxivPdfParser(Expression):
         urls = re.findall(self.url_pattern, str(data))
 
         # Convert all urls to pdf urls
-        pdf_urls = [f"https://arxiv.org/pdf/" + (f"{url.split('/')[-1]}.pdf" if 'pdf' not in url else {url.split('/')[-1]}) for url in urls]
+        pdf_urls = [
+            "https://arxiv.org/pdf/"
+            + (f"{url.split('/')[-1]}.pdf" if "pdf" not in url else {url.split("/")[-1]})
+            for url in urls
+        ]
 
         # Create temporary folder in the home directory
-        output_path = os.path.join(HOME_PATH, "temp/downloads")
-        os.makedirs(output_path, exist_ok=True)
+        output_path = HOME_PATH / "temp" / "downloads"
+        output_path.mkdir(parents=True, exist_ok=True)
 
         pdf_files = []
         with ThreadPoolExecutor() as executor:
             # Download all pdfs in parallel
-            future_to_url = {executor.submit(self.download_pdf, url, output_path): url for url in pdf_urls}
+            future_to_url = {
+                executor.submit(self.download_pdf, url, output_path): url for url in pdf_urls
+            }
             for future in as_completed(future_to_url):
                 url = future_to_url[future]
                 try:
                     pdf_files.append(future.result())
                 except Exception as exc:
-                    print('%r generated an exception: %s' % (url, exc))
+                    UserMessage(f"{url!r} generated an exception: {exc}")
 
         if len(pdf_files) == 0:
             return None
 
         # Merge all pdfs into one file
-        merged_file = self.merger(output_path, **kwargs)
+        merged_file = self.merger(str(output_path), **kwargs)
 
         # Return the merged file as a Symbol
         return_file = self._to_symbol(merged_file)
@@ -55,7 +64,7 @@ class ArxivPdfParser(Expression):
     def download_pdf(self, url, output_path):
         # Download pdfs
         response = requests.get(url)
-        file = os.path.join(output_path, f'{url.split("/")[-1]}')
-        with open(file, 'wb') as f:
+        file_path = Path(output_path) / f"{url.split('/')[-1]}"
+        with file_path.open("wb") as f:
             f.write(response.content)
-        return file
+        return str(file_path)

symai/extended/bibtex_parser.py

@@ -1,8 +1,7 @@
 from .. import core
+from ..post_processors import CodeExtractPostProcessor
 from ..pre_processors import PreProcessor
 from ..symbol import Expression, Symbol
-from ..post_processors import CodeExtractPostProcessor
-
 
 BIB_DESCRIPTION = """[Description]
 You take in a text with references to papers and return a list of biblatex entries.
@@ -69,7 +68,7 @@ Multimodal Few-Shot Learning with Frozen Language Models Maria Tsimpoukelli
 
 class BibTexPreProcessor(PreProcessor):
     def __call__(self, argument):
-        return '>>>\n{}\n\n<<<\n'.format(str(argument.args[0]))
+        return f">>>\n{argument.args[0]!s}\n\n<<<\n"
 
 
 class BibTexParser(Expression):
@@ -82,9 +81,13 @@ class BibTexParser(Expression):
         self.sym_return_type = BibTexParser
 
     def forward(self, sym: Symbol, **kwargs) -> Symbol:
-        @core.zero_shot(prompt="Create bibtex entries:\n",
-                   pre_processors=[BibTexPreProcessor()],
-                   post_processors=[CodeExtractPostProcessor()], **kwargs)
+        @core.zero_shot(
+            prompt="Create bibtex entries:\n",
+            pre_processors=[BibTexPreProcessor()],
+            post_processors=[CodeExtractPostProcessor()],
+            **kwargs,
+        )
         def _func(_, text) -> str:
             pass
+
         return _func(self, sym)

symai/extended/conversation.py

@@ -1,38 +1,43 @@
-import os
 import pickle
+from collections.abc import Callable
 from datetime import datetime
 from pathlib import Path
-from typing import Any, Callable, List, Optional
+from typing import Any
 
-from ..components import FileReader, Indexer
+from ..components import FileReader
 from ..formatter import TextContainerFormatter
 from ..interfaces import Interface
 from ..memory import SlidingWindowStringConcatMemory
 from ..symbol import Symbol
-from ..utils import CustomUserWarning, deprecated
+from ..utils import UserMessage, deprecated
 from .document import DocumentRetriever
 from .seo_query_optimizer import SEOQueryOptimizer
 
+_DEFAULT_TEXT_CONTAINER_FORMATTER = TextContainerFormatter(text_split=4)
+
 
 class CodeFormatter:
-    def __call__(self, value: str, *args: Any, **kwds: Any) -> Any:
+    def __call__(self, value: str, *_args: Any, **_kwds: Any) -> Any:
         # extract code from chat conversations or ```<language>\n{code}\n``` blocks
-        return Symbol(value).extract('Only extract code without ``` block markers or chat conversations')
+        return Symbol(value).extract(
+            "Only extract code without ``` block markers or chat conversations"
+        )
 
 
 class Conversation(SlidingWindowStringConcatMemory):
     def __init__(
-            self,
-            init: Optional[str] = None,
-            file_link: Optional[List[str]] = None,
-            url_link: Optional[List[str]] = None,
-            index_name: Optional[str] = None,
-            auto_print: bool = True,
-            truncation_percentage: float = 0.8,
-            truncation_type: str = 'head',
-            with_metadata: bool = False,
-            *args, **kwargs
-        ):
+        self,
+        init: str | None = None,
+        file_link: list[str] | None = None,
+        url_link: list[str] | None = None,
+        index_name: str | None = None,
+        auto_print: bool = True,
+        truncation_percentage: float = 0.8,
+        truncation_type: str = "head",
+        with_metadata: bool = False,
+        *args,
+        **kwargs,
+    ):
         super().__init__(*args, **kwargs)
         self.truncation_percentage = truncation_percentage
         self.truncation_type = truncation_type
@@ -46,9 +51,9 @@ class Conversation(SlidingWindowStringConcatMemory):
         self.index_name = index_name
         self.seo_opt = SEOQueryOptimizer()
         self.reader = FileReader(with_metadata=with_metadata)
-        self.scraper = Interface('naive_webscraping')
-        self.user_tag = 'USER::'
-        self.bot_tag = 'ASSISTANT::'
+        self.scraper = Interface("naive_scrape")
+        self.user_tag = "USER::"
+        self.bot_tag = "ASSISTANT::"
 
         if init is not None:
             self.store_system_message(init, *args, **kwargs)
@@ -61,14 +66,16 @@ class Conversation(SlidingWindowStringConcatMemory):
         self.indexer = None
         self.index = None
         if index_name is not None:
-            CustomUserWarning("Index not supported for conversation class.", raise_with=NotImplementedError)
+            UserMessage(
+                "Index not supported for conversation class.", raise_with=NotImplementedError
+            )
 
     def __getstate__(self):
         state = super().__getstate__().copy()
-        state.pop('seo_opt', None)
-        state.pop('indexer', None)
-        state.pop('index', None)
-        state.pop('reader', None)
+        state.pop("seo_opt", None)
+        state.pop("indexer", None)
+        state.pop("index", None)
+        state.pop("reader", None)
         return state
 
     def __setstate__(self, state):
@@ -76,41 +83,44 @@ class Conversation(SlidingWindowStringConcatMemory):
         self.seo_opt = SEOQueryOptimizer()
         self.reader = FileReader()
         if self.index_name is not None:
-            CustomUserWarning("Index not supported for conversation class.", raise_with=NotImplementedError)
+            UserMessage(
+                "Index not supported for conversation class.", raise_with=NotImplementedError
+            )
 
-    def store_system_message(self, message: str, *args, **kwargs):
-        val = f"[SYSTEM_INSTRUCTION::]: <<<\n{str(message)}\n>>>\n"
+    def store_system_message(self, message: str, *_args, **_kwargs):
+        val = f"[SYSTEM_INSTRUCTION::]: <<<\n{message!s}\n>>>\n"
         self.store(val)
 
-    def store_file(self, file_path: str, *args, **kwargs):
+    def store_file(self, file_path: str, *_args, **_kwargs):
         content = self.reader(file_path)
-        val = f"[DATA::{file_path}]: <<<\n{str(content)}\n>>>\n"
+        val = f"[DATA::{file_path}]: <<<\n{content!s}\n>>>\n"
         self.store(val)
 
-    def store_url(self, url: str, *args, **kwargs):
+    def store_url(self, url: str, *_args, **_kwargs):
         content = self.scraper(url)
-        val = f"[DATA::{url}]: <<<\n{str(content)}\n>>>\n"
+        val = f"[DATA::{url}]: <<<\n{content!s}\n>>>\n"
         self.store(val)
 
     @staticmethod
     def save_conversation_state(conversation: "Conversation", file_path: str) -> None:
         # Check if path exists and create it if it doesn't
-        dir_path = os.path.dirname(file_path)
-        os.makedirs(dir_path, exist_ok=True)
+        path_obj = Path(file_path)
+        path_obj.parent.mkdir(parents=True, exist_ok=True)
         # Save the conversation object as a pickle file
-        with open(file_path, 'wb') as handle:
+        with path_obj.open("wb") as handle:
             pickle.dump(conversation, handle, protocol=pickle.HIGHEST_PROTOCOL)
 
     def load_conversation_state(self, path: str) -> "Conversation":
         # Check if the file exists and it's not empty
-        if os.path.exists(path):
-            if os.path.getsize(path) <= 0:
-                raise Exception("File is empty.")
+        path_obj = Path(path)
+        if path_obj.exists():
+            if path_obj.stat().st_size <= 0:
+                UserMessage("File is empty.", raise_with=Exception)
             # Load the conversation object from a pickle file
-            with open(path, 'rb') as handle:
+            with path_obj.open("rb") as handle:
                 conversation_state = pickle.load(handle)
         else:
-            raise Exception("File does not exist or is empty.")
+            UserMessage("File does not exist or is empty.", raise_with=Exception)
 
         # Create a new instance of the `Conversation` class and restore
         # the state from the saved conversation
@@ -120,17 +130,19 @@ class Conversation(SlidingWindowStringConcatMemory):
         self._memory = conversation_state._memory
         self.truncation_percentage = conversation_state.truncation_percentage
         self.truncation_type = conversation_state.truncation_type
-        self.auto_print  = conversation_state.auto_print
+        self.auto_print = conversation_state.auto_print
         self.file_link = conversation_state.file_link
         self.url_link = conversation_state.url_link
-        self.index_name  = conversation_state.index_name
+        self.index_name = conversation_state.index_name
         self.seo_opt = SEOQueryOptimizer()
         self.reader = FileReader()
         if self.index_name is not None:
-            CustomUserWarning("Index not supported for conversation class.", raise_with=NotImplementedError)
+            UserMessage(
+                "Index not supported for conversation class.", raise_with=NotImplementedError
+            )
         return self
 
-    def commit(self, target_file: str = None, formatter: Optional[Callable] = None):
+    def commit(self, target_file: str | None = None, formatter: Callable | None = None):
         if target_file and isinstance(target_file, str):
             file_link = target_file
         else:
@@ -140,20 +152,22 @@ class Conversation(SlidingWindowStringConcatMemory):
             elif isinstance(file_link, list) and len(file_link) == 1:
                 file_link = file_link[0]
             else:
-                file_link = None # cannot commit to multiple files
-                raise Exception('Cannot commit to multiple files.')
+                file_link = None  # cannot commit to multiple files
+                UserMessage("Cannot commit to multiple files.", raise_with=Exception)
         if file_link:
             # if file extension is .py, then format code
             format_ = formatter
-            formatter = CodeFormatter() if format_ is None and file_link.endswith('.py') else formatter
+            formatter = (
+                CodeFormatter() if format_ is None and file_link.endswith(".py") else formatter
+            )
             val = self.value
             if formatter:
                 val = formatter(val)
             # if file does not exist, create it
-            with open(file_link, 'w') as file:
+            with Path(file_link).open("w") as file:
                 file.write(str(val))
         else:
-            raise Exception('File link is not set or a set of files.')
+            UserMessage("File link is not set or a set of files.", raise_with=Exception)
 
     def save(self, path: str, replace: bool = False) -> Symbol:
         return Symbol(self._memory).save(path, replace=replace)
@@ -161,77 +175,87 @@ class Conversation(SlidingWindowStringConcatMemory):
     def build_tag(self, tag: str, query: str) -> str:
         # get timestamp in string format
         timestamp = datetime.now().strftime("%d/%m/%Y %H:%M:%S:%f")
-        return str(f"[{tag}{timestamp}]: <<<\n{str(query)}\n>>>\n")
+        return str(f"[{tag}{timestamp}]: <<<\n{query!s}\n>>>\n")
 
     def forward(self, query: str, *args, **kwargs):
-        # dynamic takes precedence over static
-        dynamic_truncation_percentage = kwargs.get('truncation_percentage', self.truncation_percentage)
-        dynamic_truncation_type = kwargs.get('truncation_type', self.truncation_type)
-        kwargs = {**kwargs, 'truncation_percentage': dynamic_truncation_percentage, 'truncation_type': dynamic_truncation_type}
-
+        kwargs = self._apply_truncation_overrides(kwargs)
         query = self._to_symbol(query)
-        memory = None
-
-        if self.index is not None:
-            memory_split = self._memory.split(self.marker)
-            memory_shards = []
-            for ms in memory_split:
-                if ms.strip() == '':
-                    continue
-                memory_shards.append(ms)
-
-            length_memory_shards = len(memory_shards)
-            if length_memory_shards <= 3:
-                memory_shards = memory_shards
-            elif length_memory_shards <= 5:
-                memory_shards = memory_shards[:2] + memory_shards[-(length_memory_shards-2):]
-            else:
-                memory_shards = memory_shards[:2] + memory_shards[-3:]
-
-            search_query = query | '\n' | '\n'.join(memory_shards)
-            if kwargs.get('use_seo_opt'):
-                search_query = self.seo_opt(f'[Query]:' | search_query)
-            memory = self.index(search_query, *args, **kwargs)
-
-            if 'raw_result' in kwargs:
-                print(memory)
-
-        payload = ''
-        # if payload is set, then add it to the memory
-        if 'payload' in kwargs:
-            payload  = f"[Conversation Payload]:\n{kwargs.pop('payload')}\n"
-
-        index_memory = ''
-        # if index is set, then add it to the memory
-        if memory:
-            index_memory = f'[Index Retrieval]:\n{str(memory)[:1500]}\n'
-
-        payload = f'{index_memory}{payload}'
-        # perform a recall function using the query
+        memory = self._retrieve_index_memory(query, args, kwargs)
+        payload = self._build_payload(kwargs, memory)
         res = self.recall(query, *args, payload=payload, **kwargs)
 
         # if user is requesting to preview the response, then return only the preview result
-        if 'preview' in kwargs and kwargs['preview']:
+        if kwargs.get("preview"):
             if self.auto_print:
-                print(res)
+                UserMessage(str(res), style="text")
             return res
 
         ### --- asses memory update --- ###
 
-        # append the bot prompt to the memory
-        prompt = self.build_tag(self.user_tag, query)
-        self.store(prompt)
-
-        self._value = res.value # save last response
-        val = self.build_tag(self.bot_tag, res)
-        self.store(val)
+        self._append_interaction_to_memory(query, res)
 
         # WARN: DO NOT PROCESS THE RES BY REMOVING `<<<` AND `>>>` TAGS
 
         if self.auto_print:
-            print(res)
+            UserMessage(str(res), style="text")
         return res
 
+    def _apply_truncation_overrides(self, kwargs: dict[str, Any]) -> dict[str, Any]:
+        dynamic_truncation_percentage = kwargs.get(
+            "truncation_percentage", self.truncation_percentage
+        )
+        dynamic_truncation_type = kwargs.get("truncation_type", self.truncation_type)
+        return {
+            **kwargs,
+            "truncation_percentage": dynamic_truncation_percentage,
+            "truncation_type": dynamic_truncation_type,
+        }
+
+    def _retrieve_index_memory(self, query: Symbol, args: tuple[Any, ...], kwargs: dict[str, Any]):
+        if self.index is None:
+            return None
+
+        memory_split = self._memory.split(self.marker)
+        memory_shards = []
+        for shard in memory_split:
+            if shard.strip() == "":
+                continue
+            memory_shards.append(shard)
+
+        length_memory_shards = len(memory_shards)
+        if length_memory_shards > 5:
+            memory_shards = memory_shards[:2] + memory_shards[-3:]
+        elif length_memory_shards > 3:
+            retained = memory_shards[-(length_memory_shards - 2) :]
+            memory_shards = memory_shards[:2] + retained
+
+        search_query = query | "\n" | "\n".join(memory_shards)
+        if kwargs.get("use_seo_opt"):
+            search_query = self.seo_opt("[Query]:" | search_query)
+        memory = self.index(search_query, *args, **kwargs)
+
+        if "raw_result" in kwargs:
+            UserMessage(str(memory), style="text")
+        return memory
+
+    def _build_payload(self, kwargs: dict[str, Any], memory) -> str:
+        payload = ""
+        if "payload" in kwargs:
+            payload = f"[Conversation Payload]:\n{kwargs.pop('payload')}\n"
+
+        index_memory = ""
+        if memory:
+            index_memory = f"[Index Retrieval]:\n{str(memory)[:1500]}\n"
+        return f"{index_memory}{payload}"
+
+    def _append_interaction_to_memory(self, query: Symbol, res: Symbol) -> None:
+        prompt = self.build_tag(self.user_tag, query)
+        self.store(prompt)
+
+        self._value = res.value  # save last response
+        val = self.build_tag(self.bot_tag, res)
+        self.store(val)
+
 
 RETRIEVAL_CONTEXT = """[Description]
 This is a conversation between a retrieval augmented indexing program and a user. The system combines document retrieval with conversational AI to provide context-aware responses. It can:
@@ -265,27 +289,33 @@ Responses should be:
 - Referenced to source when applicable
 """
 
+
 @deprecated("Use `Conversation` instead for now. This will be removed/fixed in the future.")
 class RetrievalAugmentedConversation(Conversation):
     def __init__(
-            self,
-            folder_path: Optional[str] = None,
-            *,
-            index_name: Optional[str] = None,
-            max_depth: Optional[int] = 0,
-            auto_print: bool = True,
-            top_k: int = 5,
-            formatter: Callable = TextContainerFormatter(text_split=4),
-            overwrite: bool = False,
-            truncation_percentage: float = 0.8,
-            truncation_type: str = 'head',
-            with_metadata: bool = False,
-            raw_result: Optional[bool] = False,
-            new_dim: Optional[int] = None,
-            **kwargs
-        ):
-
-        super().__init__(auto_print=auto_print, truncation_percentage=truncation_percentage, truncation_type=truncation_type, with_metadata=with_metadata, *kwargs)
+        self,
+        folder_path: str | None = None,
+        *,
+        index_name: str | None = None,
+        max_depth: int | None = 0,
+        auto_print: bool = True,
+        top_k: int = 5,
+        formatter: Callable = _DEFAULT_TEXT_CONTAINER_FORMATTER,
+        overwrite: bool = False,
+        truncation_percentage: float = 0.8,
+        truncation_type: str = "head",
+        with_metadata: bool = False,
+        raw_result: bool | None = False,
+        new_dim: int | None = None,
+        **kwargs,
+    ):
+        super().__init__(
+            auto_print=auto_print,
+            truncation_percentage=truncation_percentage,
+            truncation_type=truncation_type,
+            with_metadata=with_metadata,
+            **kwargs,
+        )
 
         self.retriever = DocumentRetriever(
             source=folder_path,
@@ -297,7 +327,7 @@ class RetrievalAugmentedConversation(Conversation):
             with_metadata=with_metadata,
             raw_result=raw_result,
             new_dim=new_dim,
-            **kwargs
+            **kwargs,
         )
 
         self.index = self.retriever.index
@@ -322,14 +352,14 @@ class RetrievalAugmentedConversation(Conversation):
 
         memory = self.index(query, *args, **kwargs)
 
-        if 'raw_result' in kwargs:
-            print(memory)
+        if "raw_result" in kwargs:
+            UserMessage(str(memory), style="text")
             return memory
 
         prompt = self.build_tag(self.user_tag, query)
         self.store(prompt)
 
-        payload = f'[Index Retrieval]:\n{str(memory)[:1500]}\n'
+        payload = f"[Index Retrieval]:\n{str(memory)[:1500]}\n"
 
         res = self.recall(query, *args, payload=payload, **kwargs)
 
@@ -338,5 +368,5 @@ class RetrievalAugmentedConversation(Conversation):
         self.store(val)
 
         if self.auto_print:
-            print(res)
+            UserMessage(str(res), style="text")
         return res