From 3d5bbf5da977d52d9bcd6be1f09c0d20aa4e9fc6 Mon Sep 17 00:00:00 2001 From: Brandon Hancock Date: Mon, 29 Jul 2024 15:09:59 -0400 Subject: [PATCH] Fix missing rag datatype --- src/crewai/memory/storage/rag_storage.py | 11 +++++------ src/crewai/routers/pipeline_router.py | 2 +- 2 files changed, 6 insertions(+), 7 deletions(-) diff --git a/src/crewai/memory/storage/rag_storage.py b/src/crewai/memory/storage/rag_storage.py index e53f096e9..5270e9c02 100644 --- a/src/crewai/memory/storage/rag_storage.py +++ b/src/crewai/memory/storage/rag_storage.py @@ -5,12 +5,12 @@ import os import shutil from typing import Any, Dict, List, Optional -from embedchain import App -from embedchain.llm.base import BaseLlm -from embedchain.vectordb.chroma import InvalidDimensionException - from crewai.memory.storage.interface import Storage from crewai.utilities.paths import db_storage_path +from embedchain import App +from embedchain.llm.base import BaseLlm +from embedchain.models.data_type import DataType +from embedchain.vectordb.chroma import InvalidDimensionException @contextlib.contextmanager @@ -101,8 +101,7 @@ class RAGStorage(Storage): return [r for r in results if r["metadata"]["score"] >= score_threshold] def _generate_embedding(self, text: str, metadata: Dict[str, Any]) -> Any: - with suppress_logging(): - self.app.add(text, data_type="text", metadata=metadata) + self.app.add(text, data_type=DataType.TEXT, metadata=metadata) def reset(self) -> None: try: diff --git a/src/crewai/routers/pipeline_router.py b/src/crewai/routers/pipeline_router.py index d4e6cdc4b..d4b244dba 100644 --- a/src/crewai/routers/pipeline_router.py +++ b/src/crewai/routers/pipeline_router.py @@ -1,4 +1,4 @@ -from typing import Any, Callable, Dict, List, Tuple, Union +from typing import Any, Callable, Dict, Tuple, Union from pydantic import BaseModel, Field