mirror of
https://github.com/crewAIInc/crewAI.git
synced 2026-01-07 15:18:29 +00:00
14a36d3f5e
* initial knowledge * WIP * Adding core knowledge sources * Improve types and better support for file paths * added additional sources * fix linting * update yaml to include optional deps * adding in lorenze feedback * ensure embeddings are persisted * improvements all around Knowledge class * return this * properly reset memory * properly reset memory+knowledge * consolodation and improvements * linted * cleanup rm unused embedder * fix test * fix duplicate * generating cassettes for knowledge test * updated default embedder * None embedder to use default on pipeline cloning * improvements * fixed text_file_knowledge * mypysrc fixes * type check fixes * added extra cassette * just mocks * linted * mock knowledge query to not spin up db * linted * verbose run * put a flag * fix * adding docs * better docs * improvements from review * more docs * linted * rm print * more fixes * clearer docs * added docstrings and type hints for cli --------- Co-authored-by: João Moura <joaomdmoura@gmail.com> Co-authored-by: Lorenze Jay <lorenzejaytech@gmail.com>
55 lines
1.8 KiB
Python
55 lines
1.8 KiB
Python
from typing import List, Dict
|
|
from pathlib import Path
|
|
|
|
from crewai.knowledge.source.base_file_knowledge_source import BaseFileKnowledgeSource
|
|
|
|
|
|
class PDFKnowledgeSource(BaseFileKnowledgeSource):
|
|
"""A knowledge source that stores and queries PDF file content using embeddings."""
|
|
|
|
def load_content(self) -> Dict[Path, str]:
|
|
"""Load and preprocess PDF file content."""
|
|
super().load_content() # Validate the file paths
|
|
pdfplumber = self._import_pdfplumber()
|
|
|
|
paths = [self.file_path] if isinstance(self.file_path, Path) else self.file_path
|
|
content = {}
|
|
|
|
for path in paths:
|
|
text = ""
|
|
with pdfplumber.open(path) as pdf:
|
|
for page in pdf.pages:
|
|
page_text = page.extract_text()
|
|
if page_text:
|
|
text += page_text + "\n"
|
|
content[path] = text
|
|
return content
|
|
|
|
def _import_pdfplumber(self):
|
|
"""Dynamically import pdfplumber."""
|
|
try:
|
|
import pdfplumber
|
|
|
|
return pdfplumber
|
|
except ImportError:
|
|
raise ImportError(
|
|
"pdfplumber is not installed. Please install it with: pip install pdfplumber"
|
|
)
|
|
|
|
def add(self) -> None:
|
|
"""
|
|
Add PDF file content to the knowledge source, chunk it, compute embeddings,
|
|
and save the embeddings.
|
|
"""
|
|
for _, text in self.content.items():
|
|
new_chunks = self._chunk_text(text)
|
|
self.chunks.extend(new_chunks)
|
|
self.save_documents(metadata=self.metadata)
|
|
|
|
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)
|
|
]
|