Improve types and better support for file paths

src/crewai/agent.py

@@ -8,8 +8,8 @@ from pydantic import Field, InstanceOf, PrivateAttr, model_validator
 from crewai.agents import CacheHandler
 from crewai.agents.agent_builder.base_agent import BaseAgent
 from crewai.agents.crew_agent_executor import CrewAgentExecutor
-
-# from crewai.knowledge import StringKnowledgeBase
+from crewai.knowledge.knowledge import Knowledge
+from crewai.knowledge.source.base_knowledge_source import BaseKnowledgeSource
 from crewai.llm import LLM
 from crewai.memory.contextual.contextual_memory import ContextualMemory
 from crewai.tools import BaseTool
@@ -87,9 +87,9 @@ class Agent(BaseAgent):
     llm: Union[str, InstanceOf[LLM], Any] = Field(
         description="Language model that will run the agent.", default=None
     )
-    knowledge: Optional[str] = Field(
+    knowledge_sources: Optional[List[BaseKnowledgeSource]] = Field(
         default=None,
-        description="Knowledge base for the agent.",
+        description="Knowledge sources for the agent.",
     )
     function_calling_llm: Optional[Any] = Field(
         description="Language model that will run the agent.", default=None
@@ -125,6 +125,8 @@ class Agent(BaseAgent):
         default="safe",
         description="Mode for code execution: 'safe' (using Docker) or 'unsafe' (direct execution).",
     )
+    # TODO: We need to add in knowledge config (score, top_k, etc)
+    _knowledge: Optional[Knowledge] = PrivateAttr(default=None)
 
     @model_validator(mode="after")
     def post_init_setup(self):
@@ -189,7 +191,11 @@ class Agent(BaseAgent):
         if self.allow_code_execution:
             self._validate_docker_installation()
 
-        # self.knowledge = StringKnowledgeBase(content=self.knowledge)
+        # Initialize the Knowledge object if knowledge_sources are provided
+        if self.knowledge_sources:
+            self._knowledge = Knowledge(sources=self.knowledge_sources)
+        else:
+            self._knowledge = None
 
         return self
 
@@ -234,6 +240,16 @@ class Agent(BaseAgent):
             if memory.strip() != "":
                 task_prompt += self.i18n.slice("memory").format(memory=memory)
 
+        # Integrate the knowledge base
+        if self._knowledge:
+            # Query the knowledge base for relevant information
+            knowledge_snippets = self._knowledge.query(query=task.prompt())
+            print("knowledge_snippets", knowledge_snippets)
+            if knowledge_snippets:
+                formatted_knowledge = "\n".join(knowledge_snippets)
+                print("formatted_knowledge", formatted_knowledge)
+                task_prompt += f"\n\nAdditional Information:\n{formatted_knowledge}"
+
         tools = tools or self.tools or []
         self.create_agent_executor(tools=tools, task=task)
 

src/crewai/knowledge/source/base_file_knowledge_source.py

@@ -0,0 +1,24 @@
+from pathlib import Path
+
+from pydantic import Field
+
+from crewai.knowledge.source.base_knowledge_source import BaseKnowledgeSource
+
+
+class BaseFileKnowledgeSource(BaseKnowledgeSource):
+    """Base class for knowledge sources that load content from files."""
+
+    file_path: Path = Field(...)
+    content: str = Field(init=False, default="")
+
+    def model_post_init(self, context):
+        """Post-initialization method to load content."""
+        self.content = self.load_content()
+
+    def load_content(self) -> str:
+        """Load and preprocess file content. Should be overridden by subclasses."""
+        if not self.file_path.exists():
+            raise FileNotFoundError(f"File not found: {self.file_path}")
+        if not self.file_path.is_file():
+            raise ValueError(f"Path is not a file: {self.file_path}")
+        return ""

src/crewai/knowledge/source/base_knowledge_source.py

@@ -2,22 +2,20 @@ from abc import ABC, abstractmethod
 from typing import List
 
 import numpy as np
+from pydantic import BaseModel, ConfigDict, Field
 
 from crewai.knowledge.embedder.base_embedder import BaseEmbedder
 
 
-class BaseKnowledgeSource(ABC):
+class BaseKnowledgeSource(BaseModel, ABC):
     """Abstract base class for knowledge sources."""
 
-    def __init__(
-        self,
-        chunk_size: int = 1000,
-        chunk_overlap: int = 200,
-    ):
-        self.chunk_size = chunk_size
-        self.chunk_overlap = chunk_overlap
-        self.chunks: List[str] = []
-        self.chunk_embeddings: List[np.ndarray] = []
+    chunk_size: int = 1000
+    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)
 
     @abstractmethod
     def load_content(self):

src/crewai/knowledge/source/pdf_knowledge_source.py

@@ -1,22 +1,23 @@
-from pathlib import Path
 from typing import List
 
 from crewai.knowledge.embedder.base_embedder import BaseEmbedder
-from crewai.knowledge.source.base_knowledge_source import BaseKnowledgeSource
+from crewai.knowledge.source.base_file_knowledge_source import BaseFileKnowledgeSource
 
 
-class PDFKnowledgeSource(BaseKnowledgeSource):
+class PDFKnowledgeSource(BaseFileKnowledgeSource):
     """A knowledge source that stores and queries PDF file content using embeddings."""
 
-    def __init__(
-        self,
-        file_path: str,
-        chunk_size: int = 1000,
-        chunk_overlap: int = 200,
-    ):
-        super().__init__(chunk_size, chunk_overlap)
-        self.file_path = Path(file_path)
-        self.content = self.load_content()
+    def load_content(self) -> str:
+        """Load and preprocess PDF file content."""
+        super().load_content()  # Validate the file path
+        pdfplumber = self._import_pdfplumber()
+        text = ""
+        with pdfplumber.open(self.file_path) as pdf:
+            for page in pdf.pages:
+                page_text = page.extract_text()
+                if page_text:
+                    text += page_text + "\n"
+        return text
 
     def _import_pdfplumber(self):
         """Dynamically import pdfplumber."""
@@ -29,22 +30,6 @@ class PDFKnowledgeSource(BaseKnowledgeSource):
                 "pdfplumber is not installed. Please install it with: pip install pdfplumber"
             )
 
-    def load_content(self) -> str:
-        """Load and preprocess PDF file content."""
-        if not self.file_path.exists():
-            raise FileNotFoundError(f"File not found: {self.file_path}")
-        if not self.file_path.is_file():
-            raise ValueError(f"Path is not a file: {self.file_path}")
-
-        pdfplumber = self._import_pdfplumber()
-        text = ""
-        with pdfplumber.open(self.file_path) as pdf:
-            for page in pdf.pages:
-                page_text = page.extract_text()
-                if page_text:
-                    text += page_text + "\n"
-        return text
-
     def add(self, embedder: BaseEmbedder) -> None:
         """
         Add PDF file content to the knowledge source, chunk it, compute embeddings,

src/crewai/knowledge/source/string_knowledge_source.py

@@ -1,5 +1,7 @@
 from typing import List
 
+from pydantic import Field
+
 from crewai.knowledge.embedder.base_embedder import BaseEmbedder
 from crewai.knowledge.source.base_knowledge_source import BaseKnowledgeSource
 
@@ -7,29 +9,19 @@ from crewai.knowledge.source.base_knowledge_source import BaseKnowledgeSource
 class StringKnowledgeSource(BaseKnowledgeSource):
     """A knowledge source that stores and queries plain text content using embeddings."""
 
-    def __init__(
-        self,
-        content: str,
-        chunk_size: int = 1000,
-        chunk_overlap: int = 200,
-    ):
-        super().__init__(
-            chunk_size,
-            chunk_overlap,
-        )
-        self.content = content
+    content: str = Field(...)
+
+    def model_post_init(self, context):
+        """Post-initialization method to validate content."""
         self.load_content()
 
     def load_content(self):
-        """Load and preprocess string content."""
+        """Validate string content."""
         if not isinstance(self.content, str):
             raise ValueError("StringKnowledgeSource only accepts string content")
 
     def add(self, embedder: BaseEmbedder) -> None:
-        """
-        Add string content to the knowledge source, chunk it, compute embeddings,
-        and save the embeddings.
-        """
+        """Add string content to the knowledge source, chunk it, compute embeddings, and save them."""
         new_chunks = self._chunk_text(self.content)
         self.chunks.extend(new_chunks)
         # Compute embeddings for the new chunks

src/crewai/knowledge/source/text_file_knowledge_source.py

@@ -1,30 +1,15 @@
-from pathlib import Path
 from typing import List
 
 from crewai.knowledge.embedder.base_embedder import BaseEmbedder
-from crewai.knowledge.source.base_knowledge_source import BaseKnowledgeSource
+from crewai.knowledge.source.base_file_knowledge_source import BaseFileKnowledgeSource
 
 
-class TextFileKnowledgeSource(BaseKnowledgeSource):
+class TextFileKnowledgeSource(BaseFileKnowledgeSource):
     """A knowledge source that stores and queries text file content using embeddings."""
 
-    def __init__(
-        self,
-        file_path: str,
-        chunk_size: int = 1000,
-        chunk_overlap: int = 200,
-    ):
-        super().__init__(chunk_size, chunk_overlap)
-        self.file_path = Path(file_path)
-        self.content = self.load_content()
-
     def load_content(self) -> str:
         """Load and preprocess text file content."""
-        if not self.file_path.exists():
-            raise FileNotFoundError(f"File not found: {self.file_path}")
-        if not self.file_path.is_file():
-            raise ValueError(f"Path is not a file: {self.file_path}")
-
+        super().load_content()  # Validate the file path
         with self.file_path.open("r", encoding="utf-8") as f:
             return f.read()