Skip to content

Commit

Permalink
fixes mypy issues
Browse files Browse the repository at this point in the history
  • Loading branch information
Dev-Khant committed Sep 23, 2024
1 parent 7371a45 commit 66a1ecc
Show file tree
Hide file tree
Showing 5 changed files with 46 additions and 14 deletions.
12 changes: 10 additions & 2 deletions src/crewai/memory/memory.py
Original file line number Diff line number Diff line change
Expand Up @@ -23,5 +23,13 @@ def save(

self.storage.save(value, metadata)

def search(self, query: str) -> Dict[str, Any]:
return self.storage.search(query)
def search(
self,
query: str,
limit: int = 3,
filters: dict = {},
score_threshold: float = 0.35,
) -> Dict[str, Any]:
return self.storage.search(
query=query, limit=limit, filters=filters, score_threshold=score_threshold
)
12 changes: 10 additions & 2 deletions src/crewai/memory/short_term/short_term_memory.py
Original file line number Diff line number Diff line change
Expand Up @@ -42,8 +42,16 @@ def save(

super().save(value=item.data, metadata=item.metadata, agent=item.agent)

def search(self, query: str, score_threshold: float = 0.35):
return self.storage.search(query=query, score_threshold=score_threshold) # type: ignore # BUG? The reference is to the parent class, but the parent class does not have this parameters
def search(
self,
query: str,
limit: int = 3,
filters: dict = {},
score_threshold: float = 0.35,
):
return self.storage.search(
query=query, limit=limit, filters=filters, score_threshold=score_threshold
) # type: ignore # BUG? The reference is to the parent class, but the parent class does not have this parameters

def reset(self) -> None:
try:
Expand Down
4 changes: 3 additions & 1 deletion src/crewai/memory/storage/interface.py
Original file line number Diff line number Diff line change
Expand Up @@ -7,7 +7,9 @@ class Storage:
def save(self, value: Any, metadata: Dict[str, Any]) -> None:
pass

def search(self, key: str) -> Dict[str, Any]: # type: ignore
def search(
self, query: str, limit: int, filters: Dict, score_threshold: float
) -> Dict[str, Any]: # type: ignore
pass

def reset(self) -> None:
Expand Down
6 changes: 3 additions & 3 deletions src/crewai/memory/storage/rag_storage.py
Original file line number Diff line number Diff line change
Expand Up @@ -92,14 +92,14 @@ def search( # type: ignore # BUG?: Signature of "search" incompatible with supe
self,
query: str,
limit: int = 3,
filter: Optional[dict] = None,
filters: Optional[dict] = None,
score_threshold: float = 0.35,
) -> List[Any]:
with suppress_logging():
try:
results = (
self.app.search(query, limit, where=filter)
if filter
self.app.search(query, limit, where=filters)
if filters
else self.app.search(query, limit)
)
except InvalidDimensionException:
Expand Down
26 changes: 20 additions & 6 deletions src/crewai/memory/user/user_memory.py
Original file line number Diff line number Diff line change
@@ -1,5 +1,6 @@
from typing import Any, Dict, Optional

from crewai.memory.memory import Memory
from crewai.memory.user.user_memory_item import UserMemoryItem
from crewai.memory.storage.mem0_storage import Mem0Storage


Expand All @@ -15,9 +16,22 @@ def __init__(self, crew=None):
storage = Mem0Storage(type="user", crew=crew)
super().__init__(storage)

def save(self, item: UserMemoryItem) -> None:
data = f"Remember the details about the user: {item.data}"
super().save(data, item.metadata, user=item.user)
def save(
self,
value,
metadata: Optional[Dict[str, Any]] = None,
agent: Optional[str] = None,
) -> None:
data = f"Remember the details about the user: {value}"
super().save(data, metadata)

def search(self, query: str, score_threshold: float = 0.35):
return self.storage.search(query=query, score_threshold=score_threshold)
def search(
self,
query: str,
limit: int = 3,
filters: dict = {},
score_threshold: float = 0.35,
):
return super().search(
query=query, limit=limit, filters=filters, score_threshold=score_threshold
)

0 comments on commit 66a1ecc

Please sign in to comment.