Skip to content

Commit d4aa676

Browse files
feat: add configurable search parameters for RAG, knowledge, and memory (#3531)
- Add limit and score_threshold to BaseRagConfig, propagate to clients - Update default search params in RAG storage, knowledge, and memory (limit=5, threshold=0.6) - Fix linting (ruff, mypy, PERF203) and refactor save logic - Update tests for new defaults and ChromaDB behavior
1 parent 578fa8c commit d4aa676

File tree

18 files changed

+161
-106
lines changed

18 files changed

+161
-106
lines changed

src/crewai/knowledge/knowledge.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -43,7 +43,7 @@ def __init__(
4343
self.sources = sources
4444

4545
def query(
46-
self, query: list[str], results_limit: int = 3, score_threshold: float = 0.35
46+
self, query: list[str], results_limit: int = 5, score_threshold: float = 0.6
4747
) -> list[SearchResult]:
4848
"""
4949
Query across all knowledge sources to find the most relevant information.

src/crewai/knowledge/knowledge_config.py

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -9,8 +9,8 @@ class KnowledgeConfig(BaseModel):
99
score_threshold (float): The minimum score for a document to be considered relevant.
1010
"""
1111

12-
results_limit: int = Field(default=3, description="The number of results to return")
12+
results_limit: int = Field(default=5, description="The number of results to return")
1313
score_threshold: float = Field(
14-
default=0.35,
14+
default=0.6,
1515
description="The minimum score for a result to be considered relevant",
1616
)

src/crewai/knowledge/storage/base_knowledge_storage.py

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -11,9 +11,9 @@ class BaseKnowledgeStorage(ABC):
1111
def search(
1212
self,
1313
query: list[str],
14-
limit: int = 3,
14+
limit: int = 5,
1515
metadata_filter: dict[str, Any] | None = None,
16-
score_threshold: float = 0.35,
16+
score_threshold: float = 0.6,
1717
) -> list[SearchResult]:
1818
"""Search for documents in the knowledge base."""
1919

src/crewai/knowledge/storage/knowledge_storage.py

Lines changed: 9 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,5 @@
11
import logging
2+
import traceback
23
import warnings
34
from typing import Any, cast
45

@@ -49,9 +50,9 @@ def _get_client(self) -> BaseClient:
4950
def search(
5051
self,
5152
query: list[str],
52-
limit: int = 3,
53+
limit: int = 5,
5354
metadata_filter: dict[str, Any] | None = None,
54-
score_threshold: float = 0.35,
55+
score_threshold: float = 0.6,
5556
) -> list[SearchResult]:
5657
try:
5758
if not query:
@@ -73,7 +74,9 @@ def search(
7374
score_threshold=score_threshold,
7475
)
7576
except Exception as e:
76-
logging.error(f"Error during knowledge search: {e!s}")
77+
logging.error(
78+
f"Error during knowledge search: {e!s}\n{traceback.format_exc()}"
79+
)
7780
return []
7881

7982
def reset(self) -> None:
@@ -86,7 +89,9 @@ def reset(self) -> None:
8689
)
8790
client.delete_collection(collection_name=collection_name)
8891
except Exception as e:
89-
logging.error(f"Error during knowledge reset: {e!s}")
92+
logging.error(
93+
f"Error during knowledge reset: {e!s}\n{traceback.format_exc()}"
94+
)
9095

9196
def save(self, documents: list[str]) -> None:
9297
try:

src/crewai/memory/entity/entity_memory.py

Lines changed: 35 additions & 25 deletions
Original file line numberDiff line numberDiff line change
@@ -1,20 +1,20 @@
1-
from typing import Any
21
import time
2+
from typing import Any
33

44
from pydantic import PrivateAttr
55

6-
from crewai.memory.entity.entity_memory_item import EntityMemoryItem
7-
from crewai.memory.memory import Memory
8-
from crewai.memory.storage.rag_storage import RAGStorage
96
from crewai.events.event_bus import crewai_event_bus
107
from crewai.events.types.memory_events import (
11-
MemoryQueryStartedEvent,
128
MemoryQueryCompletedEvent,
139
MemoryQueryFailedEvent,
14-
MemorySaveStartedEvent,
10+
MemoryQueryStartedEvent,
1511
MemorySaveCompletedEvent,
1612
MemorySaveFailedEvent,
13+
MemorySaveStartedEvent,
1714
)
15+
from crewai.memory.entity.entity_memory_item import EntityMemoryItem
16+
from crewai.memory.memory import Memory
17+
from crewai.memory.storage.rag_storage import RAGStorage
1818

1919

2020
class EntityMemory(Memory):
@@ -31,10 +31,10 @@ def __init__(self, crew=None, embedder_config=None, storage=None, path=None):
3131
if memory_provider == "mem0":
3232
try:
3333
from crewai.memory.storage.mem0_storage import Mem0Storage
34-
except ImportError:
34+
except ImportError as e:
3535
raise ImportError(
3636
"Mem0 is not installed. Please install it with `pip install mem0ai`."
37-
)
37+
) from e
3838
config = embedder_config.get("config") if embedder_config else None
3939
storage = Mem0Storage(type="short_term", crew=crew, config=config)
4040
else:
@@ -90,23 +90,31 @@ def save(
9090
saved_count = 0
9191
errors = []
9292

93+
def save_single_item(item: EntityMemoryItem) -> tuple[bool, str | None]:
94+
"""Save a single item and return success status."""
95+
try:
96+
if self._memory_provider == "mem0":
97+
data = f"""
98+
Remember details about the following entity:
99+
Name: {item.name}
100+
Type: {item.type}
101+
Entity Description: {item.description}
102+
"""
103+
else:
104+
data = f"{item.name}({item.type}): {item.description}"
105+
106+
super(EntityMemory, self).save(data, item.metadata)
107+
return True, None
108+
except Exception as e:
109+
return False, f"{item.name}: {e!s}"
110+
93111
try:
94112
for item in items:
95-
try:
96-
if self._memory_provider == "mem0":
97-
data = f"""
98-
Remember details about the following entity:
99-
Name: {item.name}
100-
Type: {item.type}
101-
Entity Description: {item.description}
102-
"""
103-
else:
104-
data = f"{item.name}({item.type}): {item.description}"
105-
106-
super().save(data, item.metadata)
113+
success, error = save_single_item(item)
114+
if success:
107115
saved_count += 1
108-
except Exception as e:
109-
errors.append(f"{item.name}: {str(e)}")
116+
else:
117+
errors.append(error)
110118

111119
if is_batch:
112120
emit_value = f"Saved {saved_count} entities"
@@ -153,8 +161,8 @@ def save(
153161
def search(
154162
self,
155163
query: str,
156-
limit: int = 3,
157-
score_threshold: float = 0.35,
164+
limit: int = 5,
165+
score_threshold: float = 0.6,
158166
):
159167
crewai_event_bus.emit(
160168
self,
@@ -206,4 +214,6 @@ def reset(self) -> None:
206214
try:
207215
self.storage.reset()
208216
except Exception as e:
209-
raise Exception(f"An error occurred while resetting the entity memory: {e}")
217+
raise Exception(
218+
f"An error occurred while resetting the entity memory: {e}"
219+
) from e

src/crewai/memory/external/external_memory.py

Lines changed: 13 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -1,41 +1,41 @@
1-
from typing import TYPE_CHECKING, Any, Dict, Optional
21
import time
2+
from typing import TYPE_CHECKING, Any
33

4-
from crewai.memory.external.external_memory_item import ExternalMemoryItem
5-
from crewai.memory.memory import Memory
6-
from crewai.memory.storage.interface import Storage
74
from crewai.events.event_bus import crewai_event_bus
85
from crewai.events.types.memory_events import (
9-
MemoryQueryStartedEvent,
106
MemoryQueryCompletedEvent,
117
MemoryQueryFailedEvent,
12-
MemorySaveStartedEvent,
8+
MemoryQueryStartedEvent,
139
MemorySaveCompletedEvent,
1410
MemorySaveFailedEvent,
11+
MemorySaveStartedEvent,
1512
)
13+
from crewai.memory.external.external_memory_item import ExternalMemoryItem
14+
from crewai.memory.memory import Memory
15+
from crewai.memory.storage.interface import Storage
1616

1717
if TYPE_CHECKING:
1818
from crewai.memory.storage.mem0_storage import Mem0Storage
1919

2020

2121
class ExternalMemory(Memory):
22-
def __init__(self, storage: Optional[Storage] = None, **data: Any):
22+
def __init__(self, storage: Storage | None = None, **data: Any):
2323
super().__init__(storage=storage, **data)
2424

2525
@staticmethod
26-
def _configure_mem0(crew: Any, config: Dict[str, Any]) -> "Mem0Storage":
26+
def _configure_mem0(crew: Any, config: dict[str, Any]) -> "Mem0Storage":
2727
from crewai.memory.storage.mem0_storage import Mem0Storage
2828

2929
return Mem0Storage(type="external", crew=crew, config=config)
3030

3131
@staticmethod
32-
def external_supported_storages() -> Dict[str, Any]:
32+
def external_supported_storages() -> dict[str, Any]:
3333
return {
3434
"mem0": ExternalMemory._configure_mem0,
3535
}
3636

3737
@staticmethod
38-
def create_storage(crew: Any, embedder_config: Optional[Dict[str, Any]]) -> Storage:
38+
def create_storage(crew: Any, embedder_config: dict[str, Any] | None) -> Storage:
3939
if not embedder_config:
4040
raise ValueError("embedder_config is required")
4141

@@ -52,7 +52,7 @@ def create_storage(crew: Any, embedder_config: Optional[Dict[str, Any]]) -> Stor
5252
def save(
5353
self,
5454
value: Any,
55-
metadata: Optional[Dict[str, Any]] = None,
55+
metadata: dict[str, Any] | None = None,
5656
) -> None:
5757
"""Saves a value into the external storage."""
5858
crewai_event_bus.emit(
@@ -103,8 +103,8 @@ def save(
103103
def search(
104104
self,
105105
query: str,
106-
limit: int = 3,
107-
score_threshold: float = 0.35,
106+
limit: int = 5,
107+
score_threshold: float = 0.6,
108108
):
109109
crewai_event_bus.emit(
110110
self,

src/crewai/memory/memory.py

Lines changed: 7 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,4 @@
1-
from typing import Any, Dict, List, Optional, TYPE_CHECKING
1+
from typing import TYPE_CHECKING, Any, Optional
22

33
from pydantic import BaseModel
44

@@ -12,8 +12,8 @@ class Memory(BaseModel):
1212
Base class for memory, now supporting agent tags and generic metadata.
1313
"""
1414

15-
embedder_config: Optional[Dict[str, Any]] = None
16-
crew: Optional[Any] = None
15+
embedder_config: dict[str, Any] | None = None
16+
crew: Any | None = None
1717

1818
storage: Any
1919
_agent: Optional["Agent"] = None
@@ -45,7 +45,7 @@ def agent(self, agent: Optional["Agent"]) -> None:
4545
def save(
4646
self,
4747
value: Any,
48-
metadata: Optional[Dict[str, Any]] = None,
48+
metadata: dict[str, Any] | None = None,
4949
) -> None:
5050
metadata = metadata or {}
5151

@@ -54,9 +54,9 @@ def save(
5454
def search(
5555
self,
5656
query: str,
57-
limit: int = 3,
58-
score_threshold: float = 0.35,
59-
) -> List[Any]:
57+
limit: int = 5,
58+
score_threshold: float = 0.6,
59+
) -> list[Any]:
6060
return self.storage.search(
6161
query=query, limit=limit, score_threshold=score_threshold
6262
)

src/crewai/memory/short_term/short_term_memory.py

Lines changed: 13 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -1,20 +1,20 @@
1-
from typing import Any, Dict, Optional
21
import time
2+
from typing import Any
33

44
from pydantic import PrivateAttr
55

6-
from crewai.memory.memory import Memory
7-
from crewai.memory.short_term.short_term_memory_item import ShortTermMemoryItem
8-
from crewai.memory.storage.rag_storage import RAGStorage
96
from crewai.events.event_bus import crewai_event_bus
107
from crewai.events.types.memory_events import (
11-
MemoryQueryStartedEvent,
128
MemoryQueryCompletedEvent,
139
MemoryQueryFailedEvent,
14-
MemorySaveStartedEvent,
10+
MemoryQueryStartedEvent,
1511
MemorySaveCompletedEvent,
1612
MemorySaveFailedEvent,
13+
MemorySaveStartedEvent,
1714
)
15+
from crewai.memory.memory import Memory
16+
from crewai.memory.short_term.short_term_memory_item import ShortTermMemoryItem
17+
from crewai.memory.storage.rag_storage import RAGStorage
1818

1919

2020
class ShortTermMemory(Memory):
@@ -26,17 +26,17 @@ class ShortTermMemory(Memory):
2626
MemoryItem instances.
2727
"""
2828

29-
_memory_provider: Optional[str] = PrivateAttr()
29+
_memory_provider: str | None = PrivateAttr()
3030

3131
def __init__(self, crew=None, embedder_config=None, storage=None, path=None):
3232
memory_provider = embedder_config.get("provider") if embedder_config else None
3333
if memory_provider == "mem0":
3434
try:
3535
from crewai.memory.storage.mem0_storage import Mem0Storage
36-
except ImportError:
36+
except ImportError as e:
3737
raise ImportError(
3838
"Mem0 is not installed. Please install it with `pip install mem0ai`."
39-
)
39+
) from e
4040
config = embedder_config.get("config") if embedder_config else None
4141
storage = Mem0Storage(type="short_term", crew=crew, config=config)
4242
else:
@@ -56,7 +56,7 @@ def __init__(self, crew=None, embedder_config=None, storage=None, path=None):
5656
def save(
5757
self,
5858
value: Any,
59-
metadata: Optional[Dict[str, Any]] = None,
59+
metadata: dict[str, Any] | None = None,
6060
) -> None:
6161
crewai_event_bus.emit(
6262
self,
@@ -112,8 +112,8 @@ def save(
112112
def search(
113113
self,
114114
query: str,
115-
limit: int = 3,
116-
score_threshold: float = 0.35,
115+
limit: int = 5,
116+
score_threshold: float = 0.6,
117117
):
118118
crewai_event_bus.emit(
119119
self,
@@ -167,4 +167,4 @@ def reset(self) -> None:
167167
except Exception as e:
168168
raise Exception(
169169
f"An error occurred while resetting the short-term memory: {e}"
170-
)
170+
) from e

src/crewai/memory/storage/mem0_storage.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -151,7 +151,7 @@ def _last_content(messages: Iterable[dict[str, Any]], role: str) -> str:
151151
self.memory.add(conversations, **params)
152152

153153
def search(
154-
self, query: str, limit: int = 3, score_threshold: float = 0.35
154+
self, query: str, limit: int = 5, score_threshold: float = 0.6
155155
) -> list[Any]:
156156
params = {
157157
"query": query,

0 commit comments

Comments
 (0)