mirror of
https://github.com/crewAIInc/crewAI.git
synced 2026-01-09 08:08:32 +00:00
b3185ad90c
* added tool for docling support * docling support installation * use file_paths instead of file_path * fix import * organized imports * run_type docs * needs to be list * fixed logic * logged but file_path is backwards compatible * use file_paths instead of file_path 2 * added test for multiple sources for file_paths * fix run-types * enabling local files to work and type cleanup * linted * fix test and types * fixed run types * fix types * renamed to CrewDoclingSource * linted * added docs * resolve conflicts --------- Co-authored-by: Brandon Hancock (bhancock_ai) <109994880+bhancockio@users.noreply.github.com> Co-authored-by: Brandon Hancock <brandon@brandonhancock.io>
50 lines
1.6 KiB
Python
50 lines
1.6 KiB
Python
from abc import ABC, abstractmethod
|
|
from typing import Any, Dict, List, Optional
|
|
|
|
import numpy as np
|
|
from pydantic import BaseModel, ConfigDict, Field
|
|
|
|
from crewai.knowledge.storage.knowledge_storage import KnowledgeStorage
|
|
|
|
|
|
class BaseKnowledgeSource(BaseModel, ABC):
|
|
"""Abstract base class for knowledge sources."""
|
|
|
|
chunk_size: int = 4000
|
|
chunk_overlap: int = 200
|
|
chunks: List[str] = Field(default_factory=list)
|
|
chunk_embeddings: List[np.ndarray] = Field(default_factory=list)
|
|
|
|
model_config = ConfigDict(arbitrary_types_allowed=True)
|
|
storage: KnowledgeStorage = Field(default_factory=KnowledgeStorage)
|
|
metadata: Dict[str, Any] = Field(default_factory=dict) # Currently unused
|
|
collection_name: Optional[str] = Field(default=None)
|
|
|
|
@abstractmethod
|
|
def validate_content(self) -> Any:
|
|
"""Load and preprocess content from the source."""
|
|
pass
|
|
|
|
@abstractmethod
|
|
def add(self) -> None:
|
|
"""Process content, chunk it, compute embeddings, and save them."""
|
|
pass
|
|
|
|
def get_embeddings(self) -> List[np.ndarray]:
|
|
"""Return the list of embeddings for the chunks."""
|
|
return self.chunk_embeddings
|
|
|
|
def _chunk_text(self, text: str) -> List[str]:
|
|
"""Utility method to split text into chunks."""
|
|
return [
|
|
text[i : i + self.chunk_size]
|
|
for i in range(0, len(text), self.chunk_size - self.chunk_overlap)
|
|
]
|
|
|
|
def _save_documents(self):
|
|
"""
|
|
Save the documents to the storage.
|
|
This method should be called after the chunks and embeddings are generated.
|
|
"""
|
|
self.storage.save(self.chunks)
|