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

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