Knowledge (#1567)

* 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>

src/crewai/knowledge/storage/knowledge_storage.py

@@ -0,0 +1,132 @@
+import contextlib
+import io
+import logging
+import chromadb
+import os
+from crewai.utilities.paths import db_storage_path
+from typing import Optional, List
+from typing import Dict, Any
+from crewai.utilities import EmbeddingConfigurator
+from crewai.knowledge.storage.base_knowledge_storage import BaseKnowledgeStorage
+import hashlib
+
+
+@contextlib.contextmanager
+def suppress_logging(
+    logger_name="chromadb.segment.impl.vector.local_persistent_hnsw",
+    level=logging.ERROR,
+):
+    logger = logging.getLogger(logger_name)
+    original_level = logger.getEffectiveLevel()
+    logger.setLevel(level)
+    with (
+        contextlib.redirect_stdout(io.StringIO()),
+        contextlib.redirect_stderr(io.StringIO()),
+        contextlib.suppress(UserWarning),
+    ):
+        yield
+    logger.setLevel(original_level)
+
+
+class KnowledgeStorage(BaseKnowledgeStorage):
+    """
+    Extends Storage to handle embeddings for memory entries, improving
+    search efficiency.
+    """
+
+    collection: Optional[chromadb.Collection] = None
+
+    def __init__(self, embedder_config: Optional[Dict[str, Any]] = None):
+        self._initialize_app(embedder_config or {})
+
+    def search(
+        self,
+        query: List[str],
+        limit: int = 3,
+        filter: Optional[dict] = None,
+        score_threshold: float = 0.35,
+    ) -> List[Dict[str, Any]]:
+        with suppress_logging():
+            if self.collection:
+                fetched = self.collection.query(
+                    query_texts=query,
+                    n_results=limit,
+                    where=filter,
+                )
+                results = []
+                for i in range(len(fetched["ids"][0])):  # type: ignore
+                    result = {
+                        "id": fetched["ids"][0][i],  # type: ignore
+                        "metadata": fetched["metadatas"][0][i],  # type: ignore
+                        "context": fetched["documents"][0][i],  # type: ignore
+                        "score": fetched["distances"][0][i],  # type: ignore
+                    }
+                    if result["score"] >= score_threshold:  # type: ignore
+                        results.append(result)
+                return results
+            else:
+                raise Exception("Collection not initialized")
+
+    def _initialize_app(self, embedder_config: Optional[Dict[str, Any]] = None):
+        import chromadb
+        from chromadb.config import Settings
+
+        self._set_embedder_config(embedder_config)
+
+        chroma_client = chromadb.PersistentClient(
+            path=f"{db_storage_path()}/knowledge",
+            settings=Settings(allow_reset=True),
+        )
+
+        self.app = chroma_client
+
+        try:
+            self.collection = self.app.get_or_create_collection(name="knowledge")
+        except Exception:
+            raise Exception("Failed to create or get collection")
+
+    def reset(self):
+        if self.app:
+            self.app.reset()
+
+    def save(
+        self, documents: List[str], metadata: Dict[str, Any] | List[Dict[str, Any]]
+    ):
+        if self.collection:
+            metadatas = [metadata] if isinstance(metadata, dict) else metadata
+
+            ids = [
+                hashlib.sha256(doc.encode("utf-8")).hexdigest() for doc in documents
+            ]
+
+            self.collection.upsert(
+                documents=documents,
+                metadatas=metadatas,
+                ids=ids,
+            )
+        else:
+            raise Exception("Collection not initialized")
+
+    def _create_default_embedding_function(self):
+        from chromadb.utils.embedding_functions.openai_embedding_function import (
+            OpenAIEmbeddingFunction,
+        )
+
+        return OpenAIEmbeddingFunction(
+            api_key=os.getenv("OPENAI_API_KEY"), model_name="text-embedding-3-small"
+        )
+
+    def _set_embedder_config(
+        self, embedder_config: Optional[Dict[str, Any]] = None
+    ) -> None:
+        """Set the embedding configuration for the knowledge storage.
+
+        Args:
+            embedder_config (Optional[Dict[str, Any]]): Configuration dictionary for the embedder.
+                If None or empty, defaults to the default embedding function.
+        """
+        self.embedder_config = (
+            EmbeddingConfigurator().configure_embedder(embedder_config)
+            if embedder_config
+            else self._create_default_embedding_function()
+        )