diff --git a/src/crewai/knowledge/knowledge.py b/src/crewai/knowledge/knowledge.py index 4a1331799..d7729e9b7 100644 --- a/src/crewai/knowledge/knowledge.py +++ b/src/crewai/knowledge/knowledge.py @@ -1,6 +1,6 @@ import os +from typing import Any, Dict, List, Optional -from typing import List, Optional, Dict, Any from pydantic import BaseModel, ConfigDict, Field from crewai.knowledge.source.base_knowledge_source import BaseKnowledgeSource @@ -46,9 +46,7 @@ class Knowledge(BaseModel): source.storage = self.storage source.add() - def query( - self, query: List[str], limit: int = 3, preference: Optional[str] = None - ) -> List[Dict[str, Any]]: + def query(self, query: List[str], limit: int = 3) -> List[Dict[str, Any]]: """ Query across all knowledge sources to find the most relevant information. Returns the top_k most relevant chunks. @@ -57,8 +55,6 @@ class Knowledge(BaseModel): results = self.storage.search( query, limit, - filter={"preference": preference} if preference else None, - score_threshold=DEFAULT_SCORE_THRESHOLD, ) return results diff --git a/src/crewai/knowledge/source/base_file_knowledge_source.py b/src/crewai/knowledge/source/base_file_knowledge_source.py index 9cb9ec2a2..ab8929777 100644 --- a/src/crewai/knowledge/source/base_file_knowledge_source.py +++ b/src/crewai/knowledge/source/base_file_knowledge_source.py @@ -1,13 +1,13 @@ from abc import ABC, abstractmethod from pathlib import Path -from typing import Union, List, Dict, Any +from typing import Any, Dict, List, Union from pydantic import Field from crewai.knowledge.source.base_knowledge_source import BaseKnowledgeSource -from crewai.utilities.logger import Logger from crewai.knowledge.storage.knowledge_storage import KnowledgeStorage from crewai.utilities.constants import KNOWLEDGE_DIRECTORY +from crewai.utilities.logger import Logger class BaseFileKnowledgeSource(BaseKnowledgeSource, ABC): @@ -49,10 +49,9 @@ class BaseFileKnowledgeSource(BaseKnowledgeSource, ABC): color="red", ) - def save_documents(self, metadata: Dict[str, Any]): + def save_documents(self): """Save the documents to the storage.""" - chunk_metadatas = [metadata.copy() for _ in self.chunks] - self.storage.save(self.chunks, chunk_metadatas) + self.storage.save(self.chunks) def convert_to_path(self, path: Union[Path, str]) -> Path: """Convert a path to a Path object.""" diff --git a/src/crewai/knowledge/source/base_knowledge_source.py b/src/crewai/knowledge/source/base_knowledge_source.py index 6be76ca40..ddb2b9ebe 100644 --- a/src/crewai/knowledge/source/base_knowledge_source.py +++ b/src/crewai/knowledge/source/base_knowledge_source.py @@ -1,5 +1,5 @@ from abc import ABC, abstractmethod -from typing import List, Dict, Any, Optional +from typing import Any, Dict, List, Optional import numpy as np from pydantic import BaseModel, ConfigDict, Field @@ -17,7 +17,6 @@ class BaseKnowledgeSource(BaseModel, ABC): model_config = ConfigDict(arbitrary_types_allowed=True) storage: KnowledgeStorage = Field(default_factory=KnowledgeStorage) - metadata: Dict[str, Any] = Field(default_factory=dict) collection_name: Optional[str] = Field(default=None) @abstractmethod @@ -41,9 +40,9 @@ class BaseKnowledgeSource(BaseModel, ABC): for i in range(0, len(text), self.chunk_size - self.chunk_overlap) ] - def save_documents(self, metadata: Dict[str, Any]): + 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, metadata) + self.storage.save(self.chunks) diff --git a/src/crewai/knowledge/source/csv_knowledge_source.py b/src/crewai/knowledge/source/csv_knowledge_source.py index 117cbe32f..41d13af30 100644 --- a/src/crewai/knowledge/source/csv_knowledge_source.py +++ b/src/crewai/knowledge/source/csv_knowledge_source.py @@ -1,6 +1,6 @@ import csv -from typing import Dict, List from pathlib import Path +from typing import Dict, List from crewai.knowledge.source.base_file_knowledge_source import BaseFileKnowledgeSource @@ -30,7 +30,7 @@ class CSVKnowledgeSource(BaseFileKnowledgeSource): ) new_chunks = self._chunk_text(content_str) self.chunks.extend(new_chunks) - self.save_documents(metadata=self.metadata) + self.save_documents() def _chunk_text(self, text: str) -> List[str]: """Utility method to split text into chunks.""" diff --git a/src/crewai/knowledge/source/excel_knowledge_source.py b/src/crewai/knowledge/source/excel_knowledge_source.py index 5484e59d5..5656974a8 100644 --- a/src/crewai/knowledge/source/excel_knowledge_source.py +++ b/src/crewai/knowledge/source/excel_knowledge_source.py @@ -1,5 +1,6 @@ -from typing import Dict, List from pathlib import Path +from typing import Dict, List + from crewai.knowledge.source.base_file_knowledge_source import BaseFileKnowledgeSource @@ -44,7 +45,7 @@ class ExcelKnowledgeSource(BaseFileKnowledgeSource): new_chunks = self._chunk_text(content_str) self.chunks.extend(new_chunks) - self.save_documents(metadata=self.metadata) + self.save_documents() def _chunk_text(self, text: str) -> List[str]: """Utility method to split text into chunks.""" diff --git a/src/crewai/knowledge/source/json_knowledge_source.py b/src/crewai/knowledge/source/json_knowledge_source.py index ed6eec86d..e13695aa9 100644 --- a/src/crewai/knowledge/source/json_knowledge_source.py +++ b/src/crewai/knowledge/source/json_knowledge_source.py @@ -1,6 +1,6 @@ import json -from typing import Any, Dict, List from pathlib import Path +from typing import Any, Dict, List from crewai.knowledge.source.base_file_knowledge_source import BaseFileKnowledgeSource @@ -42,7 +42,7 @@ class JSONKnowledgeSource(BaseFileKnowledgeSource): ) new_chunks = self._chunk_text(content_str) self.chunks.extend(new_chunks) - self.save_documents(metadata=self.metadata) + self.save_documents() def _chunk_text(self, text: str) -> List[str]: """Utility method to split text into chunks.""" diff --git a/src/crewai/knowledge/source/pdf_knowledge_source.py b/src/crewai/knowledge/source/pdf_knowledge_source.py index e6733d78b..661a35070 100644 --- a/src/crewai/knowledge/source/pdf_knowledge_source.py +++ b/src/crewai/knowledge/source/pdf_knowledge_source.py @@ -1,5 +1,5 @@ -from typing import List, Dict from pathlib import Path +from typing import Dict, List from crewai.knowledge.source.base_file_knowledge_source import BaseFileKnowledgeSource @@ -43,7 +43,7 @@ class PDFKnowledgeSource(BaseFileKnowledgeSource): for _, text in self.content.items(): new_chunks = self._chunk_text(text) self.chunks.extend(new_chunks) - self.save_documents(metadata=self.metadata) + self.save_documents() def _chunk_text(self, text: str) -> List[str]: """Utility method to split text into chunks.""" diff --git a/src/crewai/knowledge/source/string_knowledge_source.py b/src/crewai/knowledge/source/string_knowledge_source.py index 7336fd3ea..a92706f71 100644 --- a/src/crewai/knowledge/source/string_knowledge_source.py +++ b/src/crewai/knowledge/source/string_knowledge_source.py @@ -24,7 +24,7 @@ class StringKnowledgeSource(BaseKnowledgeSource): """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) - self.save_documents(metadata=self.metadata) + self.save_documents() def _chunk_text(self, text: str) -> List[str]: """Utility method to split text into chunks.""" diff --git a/src/crewai/knowledge/source/text_file_knowledge_source.py b/src/crewai/knowledge/source/text_file_knowledge_source.py index 1ceb02843..6aa6a6ffc 100644 --- a/src/crewai/knowledge/source/text_file_knowledge_source.py +++ b/src/crewai/knowledge/source/text_file_knowledge_source.py @@ -1,5 +1,5 @@ -from typing import Dict, List from pathlib import Path +from typing import Dict, List from crewai.knowledge.source.base_file_knowledge_source import BaseFileKnowledgeSource @@ -24,7 +24,7 @@ class TextFileKnowledgeSource(BaseFileKnowledgeSource): for _, text in self.content.items(): new_chunks = self._chunk_text(text) self.chunks.extend(new_chunks) - self.save_documents(metadata=self.metadata) + self.save_documents() def _chunk_text(self, text: str) -> List[str]: """Utility method to split text into chunks.""" diff --git a/src/crewai/knowledge/storage/knowledge_storage.py b/src/crewai/knowledge/storage/knowledge_storage.py index 79aa385a8..e0e36f4a7 100644 --- a/src/crewai/knowledge/storage/knowledge_storage.py +++ b/src/crewai/knowledge/storage/knowledge_storage.py @@ -1,18 +1,20 @@ import contextlib +import hashlib import io import logging -import chromadb import os +from typing import Any, Dict, List, Optional, Union, cast +import chromadb import chromadb.errors -from crewai.utilities.paths import db_storage_path -from typing import Optional, List, Dict, Any, Union -from crewai.utilities import EmbeddingConfigurator -from crewai.knowledge.storage.base_knowledge_storage import BaseKnowledgeStorage -import hashlib -from chromadb.config import Settings from chromadb.api import ClientAPI +from chromadb.api.types import OneOrMany +from chromadb.config import Settings + +from crewai.knowledge.storage.base_knowledge_storage import BaseKnowledgeStorage +from crewai.utilities import EmbeddingConfigurator from crewai.utilities.logger import Logger +from crewai.utilities.paths import db_storage_path @contextlib.contextmanager @@ -116,11 +118,16 @@ class KnowledgeStorage(BaseKnowledgeStorage): def save( self, documents: List[str], - metadata: Union[Dict[str, Any], List[Dict[str, Any]]], + metadata: Optional[Union[Dict[str, Any], List[Dict[str, Any]]]] = None, ): if self.collection: try: - metadatas = [metadata] if isinstance(metadata, dict) else metadata + if metadata is None: + metadatas: Optional[OneOrMany[chromadb.Metadata]] = None + elif isinstance(metadata, list): + metadatas = [cast(chromadb.Metadata, m) for m in metadata] + else: + metadatas = cast(chromadb.Metadata, metadata) ids = [ hashlib.sha256(doc.encode("utf-8")).hexdigest() for doc in documents