waldiez 0.2.1__py3-none-any.whl → 0.3.0__py3-none-any.whl

waldiez/exporting/{agents → agent/utils}/rag_user/chroma_utils.py

@@ -1,3 +1,5 @@
+# SPDX-License-Identifier: Apache-2.0.
+# Copyright (c) 2024 - 2025 Waldiez and contributors.
 """Get chroma db related imports and content."""
 
 from pathlib import Path
@@ -22,19 +24,21 @@ def _get_chroma_client_string(agent: WaldiezRagUser) -> Tuple[str, str]:
         The 'client' and what to import.
     """
     # TODO: also check `connection_url` (chroma in client-server mode)
-    to_import = "chromadb"
+    to_import = "import chromadb"
     client_str = "chromadb."
     if (
         agent.retrieve_config.db_config.use_local_storage
         and agent.retrieve_config.db_config.local_storage_path is not None
     ):
-        # on windows, we get:
+        # on windows, we might get:
         # SyntaxError: (unicode error) 'unicodeescape' codec can't decode bytes
         # in position 2-3: truncated \UXXXXXXXX escape
         local_path = Path(agent.retrieve_config.db_config.local_storage_path)
         client_str += (
-            f'PersistentClient(path=r"{local_path}", '
-            "settings=Settings(anonymized_telemetry=False))"
+            "PersistentClient(\n"
+            f'    path=r"{local_path}",' + "\n"
+            "    settings=Settings(anonymized_telemetry=False),\n"
+            ")"
         )
     else:
         client_str += "Client(Settings(anonymized_telemetry=False))"
@@ -60,7 +64,7 @@ def _get_chroma_embedding_function_string(
     """
     to_import = ""
     embedding_function_arg = ""
-    embedding_function_body = ""
+    embedding_function_content = ""
     vector_db_model = agent.retrieve_config.db_config.model
     if not agent.retrieve_config.use_custom_embedding:
         to_import = (
@@ -70,13 +74,13 @@ def _get_chroma_embedding_function_string(
         embedding_function_arg = "SentenceTransformerEmbeddingFunction("
         embedding_function_arg += f'model_name="{vector_db_model}")'
     else:
-        embedding_function_arg = f"custom_embedding_function_{agent_name}"
-        embedding_function_body = (
-            f"\ndef custom_embedding_function_{agent_name}():\n"
-            f"{agent.retrieve_config.embedding_function_string}\n"
+        embedding_function_content, embedding_function_arg = (
+            agent.retrieve_config.get_custom_embedding_function(
+                name_suffix=agent_name
+            )
         )
-
-    return embedding_function_arg, to_import, embedding_function_body
+        embedding_function_content = "\n" + embedding_function_content
+    return embedding_function_arg, to_import, embedding_function_content
 
 
 def get_chroma_db_args(
@@ -108,8 +112,8 @@ def get_chroma_db_args(
     if to_import_embedding:
         to_import.add(to_import_embedding)
     kwarg_string = (
-        f"            client={agent_name}_client,\n"
-        f"            embedding_function={embedding_function_arg},\n"
+        f"            client={agent_name}_client," + "\n"
+        f"            embedding_function={embedding_function_arg}," + "\n"
     )
     # The RAG example:
     # https://ag2ai.github.io/ag2/docs/notebooks/agentchat_groupchat_RAG/
@@ -118,21 +122,22 @@ def get_chroma_db_args(
     # https://github.com/microsoft/autogen/issues/3551#issuecomment-2366930994
     # manually initializing the collection before running the flow,
     # might be a workaround.
-    content_before = f"{agent_name}_client = {client_str}\n"
+    content_before = f"{agent_name}_client = {client_str}" + "\n"
     collection_name = agent.retrieve_config.collection_name
     get_or_create = agent.retrieve_config.get_or_create
     if collection_name:
         if get_or_create:
             content_before += (
                 f"{agent_name}_client.get_or_create_collection("
-                f'"{collection_name}")\n'
+                f'"{collection_name}")' + "\n"
             )
         else:
             content_before += (
                 "try:\n"
-                f'    {agent_name}_client.get_collection("{collection_name}")\n'
+                f'    {agent_name}_client.get_collection("{collection_name}")'
+                + "\n"
                 "except ValueError:\n"
                 f"    {agent_name}_client.create_collection("
-                f'"{collection_name}")\n'
+                f'"{collection_name}")' + "\n"
             )
     return kwarg_string, to_import, embedding_function_body, content_before

waldiez/exporting/{agents → agent/utils}/rag_user/mongo_utils.py

@@ -1,3 +1,5 @@
+# SPDX-License-Identifier: Apache-2.0.
+# Copyright (c) 2024 - 2025 Waldiez and contributors.
 """Get mongodb related content and imports."""
 
 from typing import Set, Tuple
@@ -24,7 +26,7 @@ def _get_mongodb_embedding_function_string(
     """
     to_import = ""
     embedding_function_arg = ""
-    embedding_function_body = ""
+    embedding_function_content = ""
     if not agent.retrieve_config.use_custom_embedding:
         to_import = "from sentence_transformers import SentenceTransformer"
         embedding_function_arg = (
@@ -33,12 +35,13 @@ def _get_mongodb_embedding_function_string(
             ").encode"
         )
     else:
-        embedding_function_arg = f"custom_embedding_function_{agent_name}"
-        embedding_function_body = (
-            f"\ndef custom_embedding_function_{agent_name}():\n"
-            f"{agent.retrieve_config.embedding_function_string}\n"
+        embedding_function_content, embedding_function_arg = (
+            agent.retrieve_config.get_custom_embedding_function(
+                name_suffix=agent_name
+            )
         )
-    return embedding_function_arg, to_import, embedding_function_body
+        embedding_function_content = "\n" + embedding_function_content
+    return embedding_function_arg, to_import, embedding_function_content
 
 
 def get_mongodb_db_args(
@@ -67,17 +70,18 @@ def get_mongodb_db_args(
     tab = " " * 12
     db_config = agent.retrieve_config.db_config
     kwarg_string = (
-        f'{tab}connection_string="{db_config.connection_url}",\n'
-        f"{tab}embedding_function={embedding_function_arg},\n"
+        f'{tab}connection_string="{db_config.connection_url}",' + "\n"
+        f"{tab}embedding_function={embedding_function_arg}," + "\n"
     )
     wait_until_document_ready = db_config.wait_until_document_ready
     wait_until_index_ready = db_config.wait_until_index_ready
     if wait_until_document_ready is not None:
         kwarg_string += (
-            f"{tab}wait_until_document_ready={wait_until_document_ready},\n"
+            f"{tab}wait_until_document_ready={wait_until_document_ready},"
+            + "\n"
         )
     if wait_until_index_ready is not None:
         kwarg_string += (
-            f"{tab}wait_until_index_ready={wait_until_index_ready},\n"
+            f"{tab}wait_until_index_ready={wait_until_index_ready}," + "\n"
         )
     return kwarg_string, to_import, embedding_function_body

waldiez/exporting/{agents → agent/utils}/rag_user/pgvector_utils.py

@@ -1,3 +1,5 @@
+# SPDX-License-Identifier: Apache-2.0.
+# Copyright (c) 2024 - 2025 Waldiez and contributors.
 """Get pgvector related content and imports."""
 
 from typing import Set, Tuple
@@ -18,7 +20,7 @@ def _get_pgvector_client_string(agent: WaldiezRagUser) -> Tuple[str, str]:
     Tuple[str, str]
         The 'client' and what to import.
     """
-    to_import = "psycopg"
+    to_import = "import psycopg"
     client_str = "psycopg."
     connection_url = agent.retrieve_config.db_config.connection_url
     client_str += f'connect("{connection_url}")'
@@ -44,20 +46,22 @@ def _get_pgvector_embedding_function_string(
     """
     to_import = ""
     embedding_function_arg = ""
-    embedding_function_body = ""
+    embedding_function_content = ""
     if agent.retrieve_config.use_custom_embedding:
         embedding_function_arg = f"custom_embedding_function_{agent_name}"
-        embedding_function_body = (
-            f"\ndef custom_embedding_function_{agent_name}():\n"
-            f"{agent.retrieve_config.embedding_function_string}\n"
+        embedding_function_content, embedding_function_arg = (
+            agent.retrieve_config.get_custom_embedding_function(
+                name_suffix=agent_name
+            )
         )
+        embedding_function_content = "\n" + embedding_function_content
     else:
         to_import = "from sentence_transformers import SentenceTransformer"
         embedding_function_arg = "SentenceTransformer("
         embedding_function_arg += (
             f'"{agent.retrieve_config.db_config.model}").encode'
         )
-    return embedding_function_arg, to_import, embedding_function_body
+    return embedding_function_arg, to_import, embedding_function_content
 
 
 def get_pgvector_db_args(
@@ -87,7 +91,7 @@ def get_pgvector_db_args(
         else {to_import_client}
     )
     kwarg_str = (
-        f"            client={client_str},\n"
-        f"            embedding_function={embedding_function_arg},\n"
+        f"            client={client_str}," + "\n"
+        f"            embedding_function={embedding_function_arg}," + "\n"
     )
     return kwarg_str, to_import, embedding_function_body

waldiez/exporting/{agents → agent/utils}/rag_user/qdrant_utils.py

@@ -1,3 +1,5 @@
+# SPDX-License-Identifier: Apache-2.0.
+# Copyright (c) 2024 - 2025 Waldiez and contributors.
 """Get qdrant db related imports and content."""
 
 from pathlib import Path
@@ -61,7 +63,7 @@ def _get_qdrant_embedding_function_string(
     """
     to_import = ""
     embedding_function_arg = ""
-    embedding_function_body = ""
+    embedding_function_content = ""
     vector_db_model = agent.retrieve_config.db_config.model
     if not agent.retrieve_config.use_custom_embedding:
         to_import = (
@@ -71,12 +73,13 @@ def _get_qdrant_embedding_function_string(
         embedding_function_arg = "FastEmbedEmbeddingFunction("
         embedding_function_arg += f'model_name="{vector_db_model}")'
     else:
-        embedding_function_arg = f"custom_embedding_function_{agent_name}"
-        embedding_function_body = (
-            f"\ndef custom_embedding_function_{agent_name}():\n"
-            f"{agent.retrieve_config.embedding_function_string}\n"
+        embedding_function_content, embedding_function_arg = (
+            agent.retrieve_config.get_custom_embedding_function(
+                name_suffix=agent_name
+            )
         )
-    return embedding_function_arg, to_import, embedding_function_body
+        embedding_function_content = "\n" + embedding_function_content
+    return embedding_function_arg, to_import, embedding_function_content
 
 
 def get_qdrant_db_args(
@@ -106,7 +109,7 @@ def get_qdrant_db_args(
         else {to_import_client}
     )
     kwarg_string = (
-        f"            client={client_str},\n"
-        f"            embedding_function={embedding_function_arg},\n"
+        f"            client={client_str}," + "\n"
+        f"            embedding_function={embedding_function_arg}," + "\n"
     )
     return kwarg_string, to_import, embedding_function_body

waldiez/exporting/{agents → agent/utils}/rag_user/rag_user.py

@@ -1,6 +1,8 @@
+# SPDX-License-Identifier: Apache-2.0.
+# Copyright (c) 2024 - 2025 Waldiez and contributors.
 """RAG User related exporting utils."""
 
-from typing import Dict, List, Set, Tuple, Union
+from typing import Callable, Dict, List, Set, Tuple, Union
 
 from waldiez.models import (
     WaldiezAgent,
@@ -9,14 +11,63 @@ from waldiez.models import (
     WaldiezRagUserRetrieveConfig,
 )
 
-from ...utils import get_object_string, get_path_string
 from .vector_db import get_rag_user_vector_db_string
 
 
+def get_rag_user_extras(
+    agent: WaldiezAgent,
+    agent_name: str,
+    model_names: Dict[str, str],
+    path_resolver: Callable[[str], str],
+    serializer: Callable[..., str],
+) -> Tuple[str, str, Set[str]]:
+    """Get the RAG user extra argument, imports and content before the agent.
+
+    Parameters
+    ----------
+    agent : WaldiezAgent
+        The agent.
+    agent_name : str
+        The agent's name.
+    model_names : Dict[str, str]
+        A mapping from model id to model name.
+    path_resolver : Callable[[str], str]
+        The path resolver function.
+    serializer : Callable[..., str]
+        The serializer function.
+
+    Returns
+    -------
+    Tuple[str, str, Set[str]]
+        The content before the agent, the retrieve arg and the db imports.
+    """
+    before_agent_string = ""
+    retrieve_arg = ""
+    db_imports: Set[str] = set()
+    if agent.agent_type == "rag_user" and isinstance(agent, WaldiezRagUser):
+        rag_content_before_agent, retrieve_arg, db_imports = (
+            get_rag_user_retrieve_config_str(
+                agent=agent,
+                agent_name=agent_name,
+                model_names=model_names,
+                path_resolver=path_resolver,
+                serializer=serializer,
+            )
+        )
+        if retrieve_arg:
+            retrieve_arg = "\n" + f"    retrieve_config={retrieve_arg},"
+        if rag_content_before_agent:
+            before_agent_string += rag_content_before_agent
+    return before_agent_string, retrieve_arg, db_imports
+
+
+# pylint: disable=too-many-locals
 def get_rag_user_retrieve_config_str(
     agent: WaldiezRagUser,
     agent_name: str,
     model_names: Dict[str, str],
+    path_resolver: Callable[[str], str],
+    serializer: Callable[..., str],
 ) -> Tuple[str, str, Set[str]]:
     """Get the RAG user retrieve config string.
 
@@ -28,6 +79,10 @@ def get_rag_user_retrieve_config_str(
         The agent's name.
     model_names : Dict[str, str]
         A mapping from model id to model name.
+    path_resolver : Callable[[str], str]
+        The path resolver function.
+    serializer : Callable[..., str]
+        The serializer function.
     Returns
     -------
     Tuple[str, str, Set[str]]
@@ -43,32 +98,34 @@ def get_rag_user_retrieve_config_str(
         agent_name=agent_name,
     )
     imports.update(db_imports)
-    args_dict = _get_args_dict(agent, retrieve_config, model_names)
+    args_dict = _get_args_dict(
+        agent, retrieve_config, model_names, path_resolver
+    )
     if retrieve_config.use_custom_token_count:
-        token_count_arg_name = f"custom_token_count_function_{agent_name}"
-        before_the_args += (
-            f"\ndef {token_count_arg_name}():\n"
-            f"{retrieve_config.token_count_function_string}"
-            "\n\n"
+        function_content, token_count_arg_name = (
+            retrieve_config.get_custom_token_count_function(
+                name_suffix=agent_name
+            )
         )
+        before_the_args += "\n" + function_content + "\n"
         args_dict["custom_token_count_function"] = token_count_arg_name
     if retrieve_config.use_custom_text_split:
-        text_split_arg_name = f"custom_text_split_function_{agent_name}"
-        before_the_args += (
-            f"\ndef {text_split_arg_name}():\n"
-            f"{retrieve_config.text_split_function_string}"
-            "\n\n"
+        function_content, text_split_arg_name = (
+            retrieve_config.get_custom_text_split_function(
+                name_suffix=agent_name
+            )
         )
+        before_the_args += "\n" + function_content + "\n"
         args_dict["custom_text_split_function"] = text_split_arg_name
     # docs_path = args_dict.pop("docs_path", [])
-    args_content = get_object_string(args_dict)
+    args_content = serializer(args_dict)
     # get the last line (where the dict ends)
     args_parts = args_content.split("\n")
     before_vector_db = args_parts[:-1]
     closing_arg = args_parts[-1]
     args_content = "\n".join(before_vector_db)
     # add the vector_db arg
-    args_content += f',\n        "vector_db": {vector_db_arg},\n'
+    args_content += ",\n" + f'        "vector_db": {vector_db_arg},' + "\n"
     # we should not need to include the client, but let's do it
     # to avoid later issues (with telemetry or other client settings)
     # https://github.com/ag2ai/ag2/blob/main/autogen/agentchat/\
@@ -77,48 +134,11 @@ def get_rag_user_retrieve_config_str(
         f"{agent_name}_client" in before_the_args
         and agent.retrieve_config.vector_db == "chroma"
     ):
-        args_content += f'        "client": {agent_name}_client,\n'
+        args_content += f'        "client": {agent_name}_client,' + "\n"
     args_content += closing_arg
     return before_the_args, args_content, imports
 
 
-def get_rag_user_extras(
-    agent: WaldiezAgent,
-    agent_name: str,
-    model_names: Dict[str, str],
-) -> Tuple[str, str, Set[str]]:
-    """Get the RAG user extra argument, imports and content before the agent.
-
-    Parameters
-    ----------
-    agent : WaldiezAgent
-        The agent.
-    agent_name : str
-        The agent's name.
-    model_names : Dict[str, str]
-        A mapping from model id to model name.
-
-    Returns
-    -------
-    Tuple[str, str, Set[str]]
-        The content before the agent, the retrieve arg and the db imports.
-    """
-    before_agent_string = ""
-    retrieve_arg = ""
-    db_imports: Set[str] = set()
-    if agent.agent_type == "rag_user" and isinstance(agent, WaldiezRagUser):
-        rag_content_before_agent, retrieve_arg, db_imports = (
-            get_rag_user_retrieve_config_str(
-                agent=agent, agent_name=agent_name, model_names=model_names
-            )
-        )
-        if retrieve_arg:
-            retrieve_arg = f"\n    retrieve_config={retrieve_arg},"
-        if rag_content_before_agent:
-            before_agent_string += rag_content_before_agent
-    return before_agent_string, retrieve_arg, db_imports
-
-
 def _get_model_arg(
     agent: WaldiezRagUser,
     retrieve_config: WaldiezRagUserRetrieveConfig,
@@ -141,6 +161,7 @@ def _get_args_dict(
     agent: WaldiezRagUser,
     retrieve_config: WaldiezRagUserRetrieveConfig,
     model_names: Dict[str, str],
+    path_resolver: Callable[[str], str],
 ) -> Dict[str, Union[str, List[str]]]:
     model_arg = _get_model_arg(agent, retrieve_config, model_names)
     args_dict: Dict[str, Union[str, List[str]]] = {
@@ -160,12 +181,14 @@ def _get_args_dict(
             args_dict[arg] = getattr(retrieve_config, arg)
     docs_path: Union[str, List[str]] = []
     if retrieve_config.docs_path:
-        docs_path = (
+        doc_paths = (
             retrieve_config.docs_path
             if isinstance(retrieve_config.docs_path, list)
             else [retrieve_config.docs_path]
         )
-        docs_path = [get_path_string(path) for path in docs_path]
+        docs_path = [
+            item for item in [path_resolver(path) for path in doc_paths] if item
+        ]
         args_dict["docs_path"] = docs_path
     if docs_path:
         args_dict["docs_path"] = docs_path

waldiez/exporting/{agents → agent/utils}/rag_user/vector_db.py

@@ -1,7 +1,9 @@
-"""Vector DB exporting utils for RAG user agents."""
-
+# SPDX-License-Identifier: Apache-2.0.
+# Copyright (c) 2024 - 2025 Waldiez and contributors.
 # flake8: noqa E501
 # pylint: disable=line-too-long
+"""Vector DB exporting utils for RAG user agents."""
+
 from typing import Any, Set, Tuple
 
 from waldiez.models import WaldiezRagUser
@@ -31,15 +33,15 @@ def _get_metadata_arg(
     if agent.retrieve_config.db_config.metadata:
         tab = "    "
         indent = tab * 3
-        metadata_arg += f"{indent}metadata={{\n"
+        metadata_arg += f"{indent}metadata={{" + "\n"
         for key, value in agent.retrieve_config.db_config.metadata.items():
             value_string: Any = f'"{value}"'
             if str(value).isdigit():
                 value_string = int(value)
             elif str(value).replace(".", "").isdigit():
                 value_string = float(value)
-            metadata_arg += f'{indent}    "{key}": {value_string},\n'
-        metadata_arg += f"{indent}}},\n"
+            metadata_arg += f'{indent}    "{key}": {value_string},' + "\n"
+        metadata_arg += f"{indent}}}," + "\n"
     return metadata_arg
 
 
@@ -109,12 +111,12 @@ def get_rag_user_vector_db_string(
             agent, agent_name
         )
     if content_before:
-        before += f"\n{content_before}"
+        before += "\n" + f"{content_before}"
     if ef_body:
-        before += f"\n{ef_body}\n"
+        before += "\n" + f"{ef_body}" + "\n"
     if db_imports:
         imports.update(db_imports)
     kwarg_string += _get_metadata_arg(agent)
-    vdb_arg = f"{vdb_class}(\n"
+    vdb_arg = f"{vdb_class}(" + "\n"
     vdb_arg += kwarg_string + "        )"
     return before, vdb_arg, imports