waldiez 0.1.0__py3-none-any.whl

waldiez/exporting/agents/rag_user/chroma_utils.py

@@ -0,0 +1,134 @@
+"""Get chroma db related imports and content."""
+
+from pathlib import Path
+from typing import Set, Tuple
+
+from waldiez.models import WaldiezRagUser
+
+
+def _get_chroma_client_string(agent: WaldiezRagUser) -> Tuple[str, str]:
+    """Get the ChromaVectorDB client string.
+
+    Parameters
+    ----------
+    agent : WaldiezRagUser
+        The agent.
+    agent_name : str
+        The agent's name.
+
+    Returns
+    -------
+    Tuple[str, str]
+        The 'client' and what to import.
+    """
+    # TODO: also check `connection_url` (chroma in client-server mode)
+    to_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
+    ):
+        local_path = Path(agent.retrieve_config.db_config.local_storage_path)
+        client_str += f'PersistentClient(path="{local_path}")'
+    else:
+        client_str += "Client()"
+    return client_str, to_import
+
+
+def _get_chroma_embedding_function_string(
+    agent: WaldiezRagUser, agent_name: str
+) -> Tuple[str, str, str]:
+    """Get the ChromaVectorDB embedding function string.
+
+    Parameters
+    ----------
+    agent : WaldiezRagUser
+        The agent.
+    agent_name : str
+        The agent's name.
+
+    Returns
+    -------
+    Tuple[str, str, str]
+        The 'embedding_function', the import and the custom embedding function.
+    """
+    to_import = ""
+    embedding_function_arg = ""
+    embedding_function_body = ""
+    vector_db_model = agent.retrieve_config.db_config.model
+    if not agent.retrieve_config.use_custom_embedding:
+        to_import = (
+            "from chromadb.utils.embedding_functions "
+            "import SentenceTransformerEmbeddingFunction"
+        )
+        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"
+        )
+
+    return embedding_function_arg, to_import, embedding_function_body
+
+
+def get_chroma_db_args(
+    agent: WaldiezRagUser, agent_name: str
+) -> Tuple[str, Set[str], str, str]:
+    """Get the 'kwargs to use for ChromaVectorDB.
+
+    Parameters
+    ----------
+    agent : WaldiezRagUser
+        The agent.
+    agent_name : str
+        The agent's name.
+
+    Returns
+    -------
+    Tuple[str, Set[str], str]
+
+        - The 'kwargs' string.
+        - What to import.
+        - The custom embedding function.
+        - Any additional content to be used before the `kwargs` string.
+    """
+    client_str, to_import_client = _get_chroma_client_string(agent)
+    embedding_function_arg, to_import_embedding, embedding_function_body = (
+        _get_chroma_embedding_function_string(agent, agent_name)
+    )
+    to_import = {to_import_client}
+    if to_import_embedding:
+        to_import.add(to_import_embedding)
+    kwarg_string = (
+        f"            client={client_str},\n"
+        f"            embedding_function={embedding_function_arg},\n"
+    )
+    # The RAG example:
+    # https://microsoft.github.io/autogen/docs/\
+    #                                       notebooks/agentchat_groupchat_RAG
+    # raises `InvalidCollectionException`: Collection groupchat does not exist.
+    # https://github.com/chroma-core/chroma/issues/861
+    # https://github.com/microsoft/autogen/issues/3551#issuecomment-2366930994
+    # manually initializing the collection before running the flow,
+    # might be a workaround.
+    content_before = ""
+    collection_name = agent.retrieve_config.collection_name
+    get_or_create = agent.retrieve_config.get_or_create
+    if collection_name:
+        content_before = f"{agent_name}_client = {client_str}\n"
+        if get_or_create:
+            content_before += (
+                f"{agent_name}_client.get_or_create_collection("
+                f'"{collection_name}")\n'
+            )
+        else:
+            content_before += (
+                "try:\n"
+                f'    {agent_name}_client.get_collection("{collection_name}")\n'
+                "except ValueError:\n"
+                f"    {agent_name}_client.create_collection("
+                f'"{collection_name}")\n'
+            )
+    return kwarg_string, to_import, embedding_function_body, content_before

waldiez/exporting/agents/rag_user/mongo_utils.py

@@ -0,0 +1,83 @@
+"""Get mongodb related content and imports."""
+
+from typing import Set, Tuple
+
+from waldiez.models import WaldiezRagUser
+
+
+def _get_mongodb_embedding_function_string(
+    agent: WaldiezRagUser, agent_name: str
+) -> Tuple[str, str, str]:
+    """Get the MongoDBAtlasVectorDB embedding function string.
+
+    Parameters
+    ----------
+    agent : WaldiezRagUser
+        The agent.
+    agent_name : str
+        The agent's name.
+
+    Returns
+    -------
+    Tuple[str, str, str]
+        The 'embedding_function', the import and the custom_embedding_function.
+    """
+    to_import = ""
+    embedding_function_arg = ""
+    embedding_function_body = ""
+    if not agent.retrieve_config.use_custom_embedding:
+        to_import = "from sentence_transformers import SentenceTransformer"
+        embedding_function_arg = (
+            "SentenceTransformer("
+            f'"{agent.retrieve_config.db_config.model}"'
+            ").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"
+        )
+    return embedding_function_arg, to_import, embedding_function_body
+
+
+def get_mongodb_db_args(
+    agent: WaldiezRagUser, agent_name: str
+) -> Tuple[str, Set[str], str]:
+    """Get the kwargs to use for MongoDBAtlasVectorDB.
+
+    Parameters
+    ----------
+    agent : WaldiezRagUser
+        The agent.
+    agent_name : str
+        The agent's name.
+
+    Returns
+    -------
+    Tuple[str, Set[str], str]
+        The kwargs to use, what to import and the custom_embedding_function.
+    """
+    embedding_function_arg, to_import_embedding, embedding_function_body = (
+        _get_mongodb_embedding_function_string(agent, agent_name)
+    )
+    to_import: Set[str] = (
+        set() if not to_import_embedding else {to_import_embedding}
+    )
+    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"
+    )
+    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"
+        )
+    if wait_until_index_ready is not None:
+        kwarg_string += (
+            f"{tab}wait_until_index_ready={wait_until_index_ready},\n"
+        )
+    return kwarg_string, to_import, embedding_function_body

waldiez/exporting/agents/rag_user/pgvector_utils.py

@@ -0,0 +1,93 @@
+"""Get pgvector related content and imports."""
+
+from typing import Set, Tuple
+
+from waldiez.models import WaldiezRagUser
+
+
+def _get_pgvector_client_string(agent: WaldiezRagUser) -> Tuple[str, str]:
+    """Get the PGVectorDB client string.
+
+    Parameters
+    ----------
+    agent : WaldiezRagUser
+        The agent.
+
+    Returns
+    -------
+    Tuple[str, str]
+        The 'client' and what to import.
+    """
+    to_import = "psycopg"
+    client_str = "psycopg."
+    connection_url = agent.retrieve_config.db_config.connection_url
+    client_str += f'connect("{connection_url}")'
+    return client_str, to_import
+
+
+def _get_pgvector_embedding_function_string(
+    agent: WaldiezRagUser, agent_name: str
+) -> Tuple[str, str, str]:
+    """Get the PGVectorDB embedding function string.
+
+    Parameters
+    ----------
+    agent : WaldiezRagUser
+        The agent.
+    agent_name : str
+        The agent's name.
+
+    Returns
+    -------
+    Tuple[str, str, str]
+        The 'embedding_function', the import and the custom_embedding_function.
+    """
+    to_import = ""
+    embedding_function_arg = ""
+    embedding_function_body = ""
+    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"
+        )
+    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
+
+
+def get_pgvector_db_args(
+    agent: WaldiezRagUser, agent_name: str
+) -> Tuple[str, Set[str], str]:
+    """Get the kwargs to use for PGVectorDB.
+
+    Parameters
+    ----------
+    agent : WaldiezRagUser
+        The agent.
+    agent_name : str
+        The agent's name.
+
+    Returns
+    -------
+    Tuple[str, Set[str], str]
+        The kwargs to use, what to import and the custom_embedding_function.
+    """
+    client_str, to_import_client = _get_pgvector_client_string(agent)
+    embedding_function_arg, to_import_embedding, embedding_function_body = (
+        _get_pgvector_embedding_function_string(agent, agent_name)
+    )
+    to_import = (
+        {to_import_client, to_import_embedding}
+        if to_import_embedding
+        else {to_import_client}
+    )
+    kwarg_str = (
+        f"            client={client_str},\n"
+        f"            embedding_function={embedding_function_arg},\n"
+    )
+    return kwarg_str, to_import, embedding_function_body

waldiez/exporting/agents/rag_user/qdrant_utils.py

@@ -0,0 +1,112 @@
+"""Get qdrant db related imports and content."""
+
+from pathlib import Path
+from typing import Set, Tuple
+
+from waldiez.models import WaldiezRagUser
+
+
+def _get_qdrant_client_string(agent: WaldiezRagUser) -> Tuple[str, str]:
+    """Get the QdrantVectorDB client string.
+
+    Parameters
+    ----------
+    agent : WaldiezRagUser
+        The agent.
+    agent_name : str
+        The agent's name.
+
+    Returns
+    -------
+    Tuple[str, str, str]
+        The 'client' argument, and the module to import.
+    """
+    to_import: str = "from qdrant_client import QdrantClient"
+    client_str = "QdrantClient("
+    if agent.retrieve_config.db_config.use_memory:
+        client_str += 'location=":memory:")'
+    elif (
+        agent.retrieve_config.db_config.use_local_storage
+        and agent.retrieve_config.db_config.local_storage_path
+    ):
+        local_path = Path(agent.retrieve_config.db_config.local_storage_path)
+        client_str += f'location="{local_path}")'
+    elif agent.retrieve_config.db_config.connection_url:
+        client_str += (
+            f'location="{agent.retrieve_config.db_config.connection_url}")'
+        )
+    else:
+        # fallback to memory
+        client_str += 'location=":memory:")'
+    return client_str, to_import
+
+
+def _get_qdrant_embedding_function_string(
+    agent: WaldiezRagUser, agent_name: str
+) -> Tuple[str, str, str]:
+    """Get the QdrantVectorDB embedding function string.
+
+    Parameters
+    ----------
+    agent : WaldiezRagUser
+        The agent.
+    agent_name : str
+        The agent's name.
+
+    Returns
+    -------
+    Tuple[str, str, str]
+        The 'embedding_function', the module to import
+        and the custom_embedding_function if used.
+    """
+    to_import = ""
+    embedding_function_arg = ""
+    embedding_function_body = ""
+    vector_db_model = agent.retrieve_config.db_config.model
+    if not agent.retrieve_config.use_custom_embedding:
+        to_import = (
+            "from autogen.agentchat.contrib.vectordb.qdrant "
+            "import FastEmbedEmbeddingFunction"
+        )
+        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"
+        )
+    return embedding_function_arg, to_import, embedding_function_body
+
+
+def get_qdrant_db_args(
+    agent: WaldiezRagUser, agent_name: str
+) -> Tuple[str, Set[str], str]:
+    """Get the kwargs to use for QdrantVectorDB.
+
+    Parameters
+    ----------
+    agent : WaldiezRagUser
+        The agent.
+    agent_name : str
+        The agent's name.
+
+    Returns
+    -------
+    Tuple[str, Set[str], str]
+        The kwargs to use, the imports and the embedding function body if used.
+    """
+    client_str, to_import_client = _get_qdrant_client_string(agent)
+    embedding_function_arg, to_import_embedding, embedding_function_body = (
+        _get_qdrant_embedding_function_string(agent, agent_name)
+    )
+    to_import = (
+        {to_import_client, to_import_embedding}
+        if to_import_embedding
+        else {to_import_client}
+    )
+    kwarg_string = (
+        f"            client={client_str},\n"
+        f"            embedding_function={embedding_function_arg},\n"
+    )
+    return kwarg_string, to_import, embedding_function_body

waldiez/exporting/agents/rag_user/rag_user.py

@@ -0,0 +1,165 @@
+"""RAG User related exporting utils."""
+
+from typing import Dict, Set, Tuple
+
+from waldiez.models import (
+    WaldiezAgent,
+    WaldiezRagUser,
+    WaldiezRagUserModels,
+    WaldiezRagUserRetrieveConfig,
+)
+
+from ...utils import get_object_string
+from .vector_db import get_rag_user_vector_db_string
+
+
+def get_rag_user_retrieve_config_str(
+    agent: WaldiezRagUser,
+    agent_name: str,
+    model_names: Dict[str, str],
+) -> Tuple[str, str, Set[str]]:
+    """Get the RAG user retrieve config string.
+
+    Parameters
+    ----------
+    agent : WaldiezRagUser
+        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 args, the args and the imports.
+    """
+    # e.g. user_agent = RetrieveUserProxyAgent(
+    # ...other common/agent args,
+    #  retrieve_config={what_this_returns})
+    imports: Set[str] = set()
+    retrieve_config = agent.retrieve_config
+    before_the_args, vector_db_arg, db_imports = get_rag_user_vector_db_string(
+        agent=agent,
+        agent_name=agent_name,
+    )
+    imports.update(db_imports)
+    args_dict = _get_args_dict(agent, retrieve_config, model_names)
+    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"
+        )
+        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"
+        )
+        args_dict["custom_text_split_function"] = text_split_arg_name
+    args_content = get_object_string(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 += 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,
+    model_names: Dict[str, str],
+) -> str:  # pragma: no cover
+    agent_models = agent.data.model_ids
+    if agent_models:
+        first_model = agent_models[0]
+        first_model_name = model_names[first_model]
+        new_model_name = f"{first_model_name}"
+        return f"{new_model_name}"
+    if retrieve_config.model in model_names:
+        selected_model = model_names[retrieve_config.model]
+        new_model_name = f"{selected_model}"
+        return f"{new_model_name}"
+    return WaldiezRagUserModels[retrieve_config.vector_db]
+
+
+def _get_args_dict(
+    agent: WaldiezRagUser,
+    retrieve_config: WaldiezRagUserRetrieveConfig,
+    model_names: Dict[str, str],
+) -> Dict[str, str]:
+    model_arg = _get_model_arg(agent, retrieve_config, model_names)
+    args_dict = {
+        "task": retrieve_config.task,
+        "model": model_arg,
+    }
+    optional_args = [
+        "chunk_token_size",
+        "context_max_tokens",
+        "customized_prompt",
+        "customized_answer_prefix",
+        "docs_path",
+    ]
+    for arg in optional_args:
+        arg_value = getattr(retrieve_config, arg)
+        if arg_value is not None:
+            args_dict[arg] = arg_value
+            args_dict[arg] = getattr(retrieve_config, arg)
+    non_optional_args = [
+        "new_docs",
+        "update_context",
+        "get_or_create",
+        "overwrite",
+        "recursive",
+        "chunk_mode",
+        "must_break_at_empty_line",
+        "collection_name",
+        "distance_threshold",
+    ]
+    for arg in non_optional_args:
+        args_dict[arg] = getattr(retrieve_config, arg)
+    return args_dict