from typing import Union, Callable, List, Dict, Any, TypeVar


from lionagi.libs.sys_util import SysUtil

T = TypeVar("T")


def to_langchain_document(datanode: T, **kwargs: Any) -> Any:
    """
    Converts a generic data node into a Langchain Document.

    This function transforms a node, typically from another data schema, into a Langchain Document format.
    It requires the source node to have a `to_dict` method to convert it into a dictionary, then it renames specific keys
    to match the Langchain Document schema before creating a Langchain Document object.

    Args:
        datanode (T): The data node to convert. Must have a `to_dict` method.
        **kwargs: Additional keyword arguments to be passed to the Langchain Document constructor.

    Returns:
        Any: An instance of `LangchainDocument` populated with data from the input node.
    """

    SysUtil.check_import("langchain")
    from langchain.schema import Document as LangchainDocument

    dnode = datanode.to_dict()
    SysUtil.change_dict_key(dnode, old_key="content", new_key="page_content")
    SysUtil.change_dict_key(dnode, old_key="lc_id", new_key="id_")
    dnode = {**dnode, **kwargs}
    return LangchainDocument(**dnode)


def langchain_loader(
    loader: Union[str, Callable],
    loader_args: List[Any] = [],
    loader_kwargs: Dict[str, Any] = {},
) -> Any:
    """
    Initializes and uses a specified loader to load data within the Langchain ecosystem.

    This function supports dynamically selecting a loader by name or directly using a loader function.
    It passes specified arguments and keyword arguments to the loader for data retrieval or processing.

    Args:
        loader (Union[str, Callable]): A string representing the loader's name or a callable loader function.
        loader_args (List[Any], optional): A list of positional arguments for the loader.
        loader_kwargs (Dict[str, Any], optional): A dictionary of keyword arguments for the loader.

    Returns:
        Any: The result returned by the loader function, typically data loaded into a specified format.

    Raises:
        ValueError: If the loader cannot be initialized or fails to load data.

    Examples:
        >>> data = langchain_loader("json_loader", loader_args=["data.json"])
        >>> isinstance(data, dict)
        True
    """

    SysUtil.check_import("langchain")
    import langchain_community.document_loaders as document_loaders

    try:
        if isinstance(loader, str):
            loader = getattr(document_loaders, loader)
        else:
            loader = loader
    except Exception as e:
        raise ValueError(f"Invalid loader: {loader}. Error: {e}")

    try:
        loader_obj = loader(*loader_args, **loader_kwargs)
        data = loader_obj.load()
        return data
    except Exception as e:
        raise ValueError(f"Failed to load. Error: {e}")


def langchain_text_splitter(
    data: Union[str, List],
    splitter: Union[str, Callable],
    splitter_args: List[Any] = None,
    splitter_kwargs: Dict[str, Any] = None,
) -> List[str]:
    """
    Splits text or a list of texts using a specified Langchain text splitter.

    This function allows for dynamic selection of a text splitter, either by name or as a function, to split text
    or documents into chunks. The splitter can be configured with additional arguments and keyword arguments.

    Args:
        data (Union[str, List]): The text or list of texts to be split.
        splitter (Union[str, Callable]): The name of the splitter function or the splitter function itself.
        splitter_args (List[Any], optional): Positional arguments to pass to the splitter function.
        splitter_kwargs (Dict[str, Any], optional): Keyword arguments to pass to the splitter function.

    Returns:
        List[str]: A list of text chunks produced by the text splitter.

    Raises:
        ValueError: If the splitter is invalid or fails during the split operation.
    """
    splitter_args = splitter_args or []
    splitter_kwargs = splitter_kwargs or {}

    SysUtil.check_import("langchain")
    import langchain_text_splitters as text_splitter

    try:
        if isinstance(splitter, str):
            splitter = getattr(text_splitter, splitter)
        else:
            splitter = splitter
    except Exception as e:
        raise ValueError(f"Invalid text splitter: {splitter}. Error: {e}")

    try:
        splitter_obj = splitter(*splitter_args, **splitter_kwargs)
        if isinstance(data, str):
            chunk = splitter_obj.split_text(data)
        else:
            chunk = splitter_obj.split_documents(data)
        return chunk
    except Exception as e:
        raise ValueError(f"Failed to split. Error: {e}")

langchain.schema.Document
[((110, 122), 'typing.TypeVar', 'TypeVar', (['"""T"""'], {}), "('T')\n", (117, 122), False, 'from typing import Union, Callable, List, Dict, Any, TypeVar\n'), ((877, 910), 'lionagi.libs.sys_util.SysUtil.check_import', 'SysUtil.check_import', (['"""langchain"""'], {}), "('langchain')\n", (897, 910), False, 'from lionagi.libs.sys_util import SysUtil\n'), ((1010, 1083), 'lionagi.libs.sys_util.SysUtil.change_dict_key', 'SysUtil.change_dict_key', (['dnode'], {'old_key': '"""content"""', 'new_key': '"""page_content"""'}), "(dnode, old_key='content', new_key='page_content')\n", (1033, 1083), False, 'from lionagi.libs.sys_util import SysUtil\n'), ((1088, 1150), 'lionagi.libs.sys_util.SysUtil.change_dict_key', 'SysUtil.change_dict_key', (['dnode'], {'old_key': '"""lc_id"""', 'new_key': '"""id_"""'}), "(dnode, old_key='lc_id', new_key='id_')\n", (1111, 1150), False, 'from lionagi.libs.sys_util import SysUtil\n'), ((1194, 1220), 'langchain.schema.Document', 'LangchainDocument', ([], {}), '(**dnode)\n', (1211, 1220), True, 'from langchain.schema import Document as LangchainDocument\n'), ((2342, 2375), 'lionagi.libs.sys_util.SysUtil.check_import', 'SysUtil.check_import', (['"""langchain"""'], {}), "('langchain')\n", (2362, 2375), False, 'from lionagi.libs.sys_util import SysUtil\n'), ((4063, 4096), 'lionagi.libs.sys_util.SysUtil.check_import', 'SysUtil.check_import', (['"""langchain"""'], {}), "('langchain')\n", (4083, 4096), False, 'from lionagi.libs.sys_util import SysUtil\n')]