feat: replace embedchain with native crewai adapter

crewai_tools/adapters/crewai_rag_adapter.py

@@ -0,0 +1,215 @@
+"""Adapter for CrewAI's native RAG system."""
+
+from typing import Any, TypedDict, TypeAlias
+from typing_extensions import Unpack
+from pathlib import Path
+import hashlib
+
+from pydantic import Field, PrivateAttr
+from crewai.rag.config.utils import get_rag_client
+from crewai.rag.config.types import RagConfigType
+from crewai.rag.types import BaseRecord, SearchResult
+from crewai.rag.core.base_client import BaseClient
+from crewai.rag.factory import create_client
+
+from crewai_tools.tools.rag.rag_tool import Adapter
+from crewai_tools.rag.data_types import DataType
+from crewai_tools.rag.misc import sanitize_metadata_for_chromadb
+from crewai_tools.rag.chunkers.base_chunker import BaseChunker
+
+ContentItem: TypeAlias = str | Path | dict[str, Any]
+
+class AddDocumentParams(TypedDict, total=False):
+    """Parameters for adding documents to the RAG system."""
+    data_type: DataType
+    metadata: dict[str, Any]
+    website: str
+    url: str
+    file_path: str | Path
+    github_url: str
+    youtube_url: str
+    directory_path: str | Path
+
+
+class CrewAIRagAdapter(Adapter):
+    """Adapter that uses CrewAI's native RAG system.
+
+    Supports custom vector database configuration through the config parameter.
+    """
+
+    collection_name: str = "default"
+    summarize: bool = False
+    similarity_threshold: float = 0.6
+    limit: int = 5
+    config: RagConfigType | None = None
+    _client: BaseClient | None = PrivateAttr(default=None)
+
+    def model_post_init(self, __context: Any) -> None:
+        """Initialize the CrewAI RAG client after model initialization."""
+        if self.config is not None:
+            self._client = create_client(self.config)
+        else:
+            self._client = get_rag_client()
+        self._client.get_or_create_collection(collection_name=self.collection_name)
+
+    def query(self, question: str, similarity_threshold: float | None = None, limit: int | None = None) -> str:
+        """Query the knowledge base with a question.
+
+        Args:
+            question: The question to ask
+            similarity_threshold: Minimum similarity score for results (default: 0.6)
+            limit: Maximum number of results to return (default: 5)
+
+        Returns:
+            Relevant content from the knowledge base
+        """
+        search_limit = limit if limit is not None else self.limit
+        search_threshold = similarity_threshold if similarity_threshold is not None else self.similarity_threshold
+
+        results: list[SearchResult] = self._client.search(
+            collection_name=self.collection_name,
+            query=question,
+            limit=search_limit,
+            score_threshold=search_threshold
+        )
+
+        if not results:
+            return "No relevant content found."
+
+        contents: list[str] = []
+        for result in results:
+            content: str = result.get("content", "")
+            if content:
+                contents.append(content)
+
+        return "\n\n".join(contents)
+
+    def add(self, *args: ContentItem, **kwargs: Unpack[AddDocumentParams]) -> None:
+        """Add content to the knowledge base.
+
+        This method handles various input types and converts them to documents
+        for the vector database. It supports the data_type parameter for 
+        compatibility with existing tools.
+
+        Args:
+            *args: Content items to add (strings, paths, or document dicts)
+            **kwargs: Additional parameters including data_type, metadata, etc.
+        """
+        from crewai_tools.rag.data_types import DataTypes, DataType
+        from crewai_tools.rag.source_content import SourceContent
+        from crewai_tools.rag.base_loader import LoaderResult
+        import os
+
+        documents: list[BaseRecord] = []
+        data_type: DataType | None = kwargs.get("data_type")
+        base_metadata: dict[str, Any] = kwargs.get("metadata", {})
+
+        for arg in args:
+            source_ref: str
+            if isinstance(arg, dict):
+                source_ref = str(arg.get("source", arg.get("content", "")))
+            else:
+                source_ref = str(arg)
+
+            if not data_type:
+                data_type = DataTypes.from_content(source_ref)
+
+            if data_type == DataType.DIRECTORY:
+                if not os.path.isdir(source_ref):
+                    raise ValueError(f"Directory does not exist: {source_ref}")
+
+                # Define binary and non-text file extensions to skip
+                binary_extensions = {'.pyc', '.pyo', '.png', '.jpg', '.jpeg', '.gif', 
+                                    '.bmp', '.ico', '.svg', '.webp', '.pdf', '.zip', 
+                                    '.tar', '.gz', '.bz2', '.7z', '.rar', '.exe', 
+                                    '.dll', '.so', '.dylib', '.bin', '.dat', '.db',
+                                    '.sqlite', '.class', '.jar', '.war', '.ear'}
+
+                for root, dirs, files in os.walk(source_ref):
+                    dirs[:] = [d for d in dirs if not d.startswith('.')]
+
+                    for filename in files:
+                        if filename.startswith('.'):
+                            continue
+
+                        # Skip binary files based on extension
+                        file_ext = os.path.splitext(filename)[1].lower()
+                        if file_ext in binary_extensions:
+                            continue
+
+                        # Skip __pycache__ directories
+                        if '__pycache__' in root:
+                            continue
+
+                        file_path: str = os.path.join(root, filename)
+                        try:
+                            file_data_type: DataType = DataTypes.from_content(file_path)
+                            file_loader = file_data_type.get_loader()
+                            file_chunker = file_data_type.get_chunker()
+
+                            file_source = SourceContent(file_path)
+                            file_result: LoaderResult = file_loader.load(file_source)
+
+                            file_chunks = file_chunker.chunk(file_result.content)
+
+                            for chunk_idx, file_chunk in enumerate(file_chunks):
+                                file_metadata: dict[str, Any] = base_metadata.copy()
+                                file_metadata.update(file_result.metadata)
+                                file_metadata["data_type"] = str(file_data_type)
+                                file_metadata["file_path"] = file_path
+                                file_metadata["chunk_index"] = chunk_idx
+                                file_metadata["total_chunks"] = len(file_chunks)
+
+                                if isinstance(arg, dict):
+                                    file_metadata.update(arg.get("metadata", {}))
+
+                                chunk_id = hashlib.sha256(f"{file_result.doc_id}_{chunk_idx}_{file_chunk}".encode()).hexdigest()
+
+                                documents.append({
+                                    "doc_id": chunk_id,
+                                    "content": file_chunk,
+                                    "metadata": sanitize_metadata_for_chromadb(file_metadata)
+                                })
+                        except Exception:
+                            # Silently skip files that can't be processed
+                            continue
+            else:
+                metadata: dict[str, Any] = base_metadata.copy()
+
+                if data_type in [DataType.PDF_FILE, DataType.TEXT_FILE, DataType.DOCX, 
+                                  DataType.CSV, DataType.JSON, DataType.XML, DataType.MDX]:
+                    if not os.path.isfile(source_ref):
+                        raise FileNotFoundError(f"File does not exist: {source_ref}")
+
+                loader = data_type.get_loader()
+                chunker = data_type.get_chunker()
+
+                source_content = SourceContent(source_ref)
+                loader_result: LoaderResult = loader.load(source_content)
+
+                chunks = chunker.chunk(loader_result.content)
+
+                for i, chunk in enumerate(chunks):
+                    chunk_metadata: dict[str, Any] = metadata.copy()
+                    chunk_metadata.update(loader_result.metadata)
+                    chunk_metadata["data_type"] = str(data_type)
+                    chunk_metadata["chunk_index"] = i
+                    chunk_metadata["total_chunks"] = len(chunks)
+                    chunk_metadata["source"] = source_ref
+
+                    if isinstance(arg, dict):
+                        chunk_metadata.update(arg.get("metadata", {}))
+
+                    chunk_id = hashlib.sha256(f"{loader_result.doc_id}_{i}_{chunk}".encode()).hexdigest()
+
+                    documents.append({
+                        "doc_id": chunk_id,
+                        "content": chunk,
+                        "metadata": sanitize_metadata_for_chromadb(chunk_metadata)
+                    })
+
+        if documents:
+            self._client.add_documents(
+                collection_name=self.collection_name,
+                documents=documents
+            )

crewai_tools/rag/chunkers/base_chunker.py

@@ -112,7 +112,10 @@ def _merge_splits(self, splits: List[str], separator: str) -> List[str]:
                 if separator == "":
                     doc = "".join(current_doc)
                 else:
-                    doc = separator.join(current_doc)
+                    if self._keep_separator and separator == " ":
+                        doc = "".join(current_doc)
+                    else:
+                        doc = separator.join(current_doc)
 
                 if doc:
                     docs.append(doc)
@@ -133,7 +136,10 @@ def _merge_splits(self, splits: List[str], separator: str) -> List[str]:
             if separator == "":
                 doc = "".join(current_doc)
             else:
-                doc = separator.join(current_doc)
+                if self._keep_separator and separator == " ":
+                    doc = "".join(current_doc)
+                else:
+                    doc = separator.join(current_doc)
 
             if doc:
                 docs.append(doc)

crewai_tools/rag/data_types.py

@@ -25,6 +25,8 @@ class DataType(str, Enum):
     # Web types
     WEBSITE = "website"
     DOCS_SITE = "docs_site"
+    YOUTUBE_VIDEO = "youtube_video"
+    YOUTUBE_CHANNEL = "youtube_channel"
 
     # Raw types
     TEXT = "text"
@@ -34,6 +36,7 @@ def get_chunker(self) -> BaseChunker:
         from importlib import import_module
 
         chunkers = {
+            DataType.PDF_FILE: ("text_chunker", "TextChunker"),
             DataType.TEXT_FILE: ("text_chunker", "TextChunker"),
             DataType.TEXT: ("text_chunker", "TextChunker"),
             DataType.DOCX: ("text_chunker", "DocxChunker"),
@@ -45,9 +48,18 @@ def get_chunker(self) -> BaseChunker:
             DataType.XML: ("structured_chunker", "XmlChunker"),
 
             DataType.WEBSITE: ("web_chunker", "WebsiteChunker"),
+            DataType.DIRECTORY: ("text_chunker", "TextChunker"),
+            DataType.YOUTUBE_VIDEO: ("text_chunker", "TextChunker"),
+            DataType.YOUTUBE_CHANNEL: ("text_chunker", "TextChunker"),
+            DataType.GITHUB: ("text_chunker", "TextChunker"),
+            DataType.DOCS_SITE: ("text_chunker", "TextChunker"),
+            DataType.MYSQL: ("text_chunker", "TextChunker"),
+            DataType.POSTGRES: ("text_chunker", "TextChunker"),
         }
 
-        module_name, class_name = chunkers.get(self, ("default_chunker", "DefaultChunker"))
+        if self not in chunkers:
+            raise ValueError(f"No chunker defined for {self}")
+        module_name, class_name = chunkers[self]
         module_path = f"crewai_tools.rag.chunkers.{module_name}"
 
         try:
@@ -60,6 +72,7 @@ def get_loader(self) -> BaseLoader:
         from importlib import import_module
 
         loaders = {
+            DataType.PDF_FILE: ("pdf_loader", "PDFLoader"),
             DataType.TEXT_FILE: ("text_loader", "TextFileLoader"),
             DataType.TEXT: ("text_loader", "TextLoader"),
             DataType.XML: ("xml_loader", "XMLLoader"),
@@ -69,9 +82,17 @@ def get_loader(self) -> BaseLoader:
             DataType.DOCX: ("docx_loader", "DOCXLoader"),
             DataType.CSV: ("csv_loader", "CSVLoader"),
             DataType.DIRECTORY: ("directory_loader", "DirectoryLoader"),
+            DataType.YOUTUBE_VIDEO: ("youtube_video_loader", "YoutubeVideoLoader"),
+            DataType.YOUTUBE_CHANNEL: ("youtube_channel_loader", "YoutubeChannelLoader"),
+            DataType.GITHUB: ("github_loader", "GithubLoader"),
+            DataType.DOCS_SITE: ("docs_site_loader", "DocsSiteLoader"),
+            DataType.MYSQL: ("mysql_loader", "MySQLLoader"),
+            DataType.POSTGRES: ("postgres_loader", "PostgresLoader"),
         }
 
-        module_name, class_name = loaders.get(self, ("text_loader", "TextLoader"))
+        if self not in loaders:
+            raise ValueError(f"No loader defined for {self}")
+        module_name, class_name = loaders[self]
         module_path = f"crewai_tools.rag.loaders.{module_name}"
         try:
             module = import_module(module_path)

crewai_tools/rag/loaders/__init__.py

@@ -6,6 +6,9 @@
 from crewai_tools.rag.loaders.docx_loader import DOCXLoader
 from crewai_tools.rag.loaders.csv_loader import CSVLoader
 from crewai_tools.rag.loaders.directory_loader import DirectoryLoader
+from crewai_tools.rag.loaders.pdf_loader import PDFLoader
+from crewai_tools.rag.loaders.youtube_video_loader import YoutubeVideoLoader
+from crewai_tools.rag.loaders.youtube_channel_loader import YoutubeChannelLoader
 
 __all__ = [
     "TextFileLoader",
@@ -17,4 +20,7 @@
     "DOCXLoader",
     "CSVLoader",
     "DirectoryLoader",
+    "PDFLoader",
+    "YoutubeVideoLoader",
+    "YoutubeChannelLoader",
 ]