mem0ai-azure-mysql 0.1.116.3__tar.gz → 0.1.116.4__tar.gz

{mem0ai_azure_mysql-0.1.116.3 → mem0ai_azure_mysql-0.1.116.4}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: mem0ai-azure-mysql
-Version: 0.1.116.3
+Version: 0.1.116.4
 Summary: Long-term memory for AI Agents with Azure DefaultAzureCredential authentication and MySQL history database support
 Requires-Python: <4.0,>=3.9
 Requires-Dist: azure-identity>=1.23.1

{mem0ai_azure_mysql-0.1.116.3 → mem0ai_azure_mysql-0.1.116.4}/mem0/embeddings/azure_openai.py

@@ -6,6 +6,7 @@ from azure.identity import DefaultAzureCredential, get_bearer_token_provider
 
 from mem0.configs.embeddings.base import BaseEmbedderConfig
 from mem0.embeddings.base import EmbeddingBase
+from mem0.memory.utils import time_perf
 
 
 class AzureOpenAIEmbedding(EmbeddingBase):
@@ -41,6 +42,7 @@ class AzureOpenAIEmbedding(EmbeddingBase):
             **auth_kwargs,
         )
 
+    @time_perf
     def embed(self, text, memory_action: Optional[Literal["add", "search", "update"]] = None):
         """
         Get the embedding for the given text using OpenAI.
@@ -53,3 +55,18 @@ class AzureOpenAIEmbedding(EmbeddingBase):
         """
         text = text.replace("\n", " ")
         return self.client.embeddings.create(input=[text], model=self.config.model).data[0].embedding
+
+    @time_perf
+    def embed_in_batch(self, texts, memory_action: Optional[Literal["add", "search", "update"]] = None):
+        """
+        Get the embedding for the given texts.
+
+        Args:
+            texts (list[str]): The text list to embed.
+            memory_action (optional): The type of embedding to use. Must be one of "add", "search", or "update". Defaults to None.
+        Returns:
+            list: List of the embedding vector.
+        """
+        texts = [text.replace("\n", " ") for text in texts]
+        resp = self.client.embeddings.create(input=texts, model=self.config.model)
+        return [item.embedding for item in resp.data]

{mem0ai_azure_mysql-0.1.116.3 → mem0ai_azure_mysql-0.1.116.4}/mem0/embeddings/base.py

@@ -18,7 +18,7 @@ class EmbeddingBase(ABC):
             self.config = config
 
     @abstractmethod
-    def embed(self, text, memory_action: Optional[Literal["add", "search", "update"]]):
+    def embed(self, text, memory_action: Optional[Literal["add", "search", "update"]]) -> list[float]:
         """
         Get the embedding for the given text.
 
@@ -29,3 +29,15 @@ class EmbeddingBase(ABC):
             list: The embedding vector.
         """
         pass
+
+    def embed_in_batch(self, texts, memory_action: Optional[Literal["add", "search", "update"]]) -> list[list[float]]:
+        """
+        Get the embedding for the given texts.
+
+        Args:
+            texts (list[str]): The text list to embed.
+            memory_action (optional): The type of embedding to use. Must be one of "add", "search", or "update". Defaults to None.
+        Returns:
+            list: List of the embedding vector.
+        """
+        return [self.embed(text, memory_action) for text in texts]

{mem0ai_azure_mysql-0.1.116.3 → mem0ai_azure_mysql-0.1.116.4}/mem0/memory/graph_memory.py

@@ -3,7 +3,7 @@ import time
 
 import cohere
 
-from mem0.memory.utils import format_entities, sanitize_relationship_for_cypher
+from mem0.memory.utils import format_entities, sanitize_relationship_for_cypher, time_perf
 
 try:
     from langchain_neo4j import Neo4jGraph
@@ -114,8 +114,6 @@ class MemoryGraph:
         logger.info(f"Searching for query in neo4j: {query} with filters: {filters}")
         entity_type_map = self._retrieve_nodes_from_data(query, filters)
         search_output = self._search_graph_db(node_list=list(entity_type_map.keys()), filters=filters)
-        print(f'entity_type_map: {entity_type_map}')
-        print(f'search_output: {search_output}')
 
         if not search_output:
             return []
@@ -158,14 +156,12 @@ class MemoryGraph:
             base_url=rerank_config.get("base_url"),
         )
         docs = [f'{item["source"]} {item["relationship"]} {item["destination"]}' for item in search_output]
-        print(f'docs: {docs}')
         response = co.rerank(
             model=rerank_config.get("model"),
             query=query,
             documents=docs,
             top_n=self.top_k,
         )
-        print(f'response.results: {response.results}')
 
         reranked_results = []
         for result in response.results:
@@ -176,7 +172,6 @@ class MemoryGraph:
                 original_output["destination"],
                 result.relevance_score
             ])
-        print(f'reranked_results: {reranked_results}')
         return reranked_results
 
     def delete_all(self, filters):
@@ -243,6 +238,7 @@ class MemoryGraph:
 
         return final_results
 
+    @time_perf
     def _retrieve_nodes_from_data(self, data, filters):
         """Extracts all the entities mentioned in the query."""
         _tools = [EXTRACT_ENTITIES_TOOL]
@@ -276,6 +272,7 @@ class MemoryGraph:
         logger.debug(f"Entity type map: {entity_type_map}\n search_results={search_results}")
         return entity_type_map
 
+    @time_perf
     def _establish_nodes_relations_from_data(self, data, filters, entity_type_map):
         """Establish relations among the extracted nodes."""
 
@@ -318,6 +315,7 @@ class MemoryGraph:
         logger.debug(f"Extracted entities: {entities}")
         return entities
 
+    @time_perf
     def _search_graph_db(self, node_list, filters, limit=100):
         """Search similar nodes among and their respective incoming and outgoing relations."""
         result_relations = []
@@ -329,9 +327,9 @@ class MemoryGraph:
         if filters.get("run_id"):
             node_props.append("run_id: $run_id")
         node_props_str = ", ".join(node_props)
+        n_embeddings = self.embedding_model.embed_in_batch(node_list)
 
-        for node in node_list:
-            n_embedding = self.embedding_model.embed(node)
+        for n_embedding in n_embeddings:
 
             cypher_query = f"""
             MATCH (n {self.node_label} {{{node_props_str}}})
@@ -377,7 +375,8 @@ class MemoryGraph:
                 unique_relations.append(rel)
 
         return unique_relations
-
+   
+    @time_perf
     def _get_delete_entities_from_search_output(self, search_output, data, filters):
         """Get the entities to be deleted from the search output."""
         search_output_string = format_entities(search_output)
@@ -414,6 +413,7 @@ class MemoryGraph:
         logger.debug(f"Deleted relationships: {to_be_deleted}")
         return to_be_deleted
 
+    @time_perf
     def _delete_entities(self, to_be_deleted, filters):
         """Delete the entities from the graph."""
         user_id = filters["user_id"]
@@ -469,6 +469,7 @@ class MemoryGraph:
 
         return results
 
+    @time_perf
     def _add_entities(self, to_be_added, filters, entity_type_map):
         """Add the new entities to the graph. Merge the nodes if they already exist."""
         user_id = filters["user_id"]

{mem0ai_azure_mysql-0.1.116.3 → mem0ai_azure_mysql-0.1.116.4}/mem0/memory/main.py

@@ -29,6 +29,7 @@ from mem0.memory.utils import (
     parse_vision_messages,
     process_telemetry_filters,
     remove_code_blocks,
+    time_perf,
 )
 from mem0.utils.factory import (
     EmbedderFactory,
@@ -282,6 +283,7 @@ class Memory(MemoryBase):
 
         return {"results": vector_store_result}
 
+    @time_perf
     def _add_to_vector_store(self, messages, metadata, filters, infer):
         if not infer:
             returned_memories = []
@@ -451,6 +453,7 @@ class Memory(MemoryBase):
         )
         return returned_memories
 
+    @time_perf
     def _add_to_graph(self, messages, filters):
         added_entities = []
         if self.enable_graph:

{mem0ai_azure_mysql-0.1.116.3 → mem0ai_azure_mysql-0.1.116.4}/mem0/memory/utils.py

@@ -1,8 +1,12 @@
-import hashlib
 import re
+import time
+import hashlib
+import logging
+from functools import wraps
 
 from mem0.configs.prompts import FACT_RETRIEVAL_PROMPT
 
+logger = logging.getLogger(__name__)
 
 def get_fact_retrieval_messages(message):
     return FACT_RETRIEVAL_PROMPT, f"Input:\n{message}"
@@ -182,3 +186,15 @@ def sanitize_relationship_for_cypher(relationship) -> str:
         sanitized = sanitized.replace(old, new)
 
     return re.sub(r"_+", "_", sanitized).strip("_")
+
+
+def time_perf(func):
+    @wraps(func)
+    def wrapper(*args, **kwargs):
+        start = time.perf_counter()
+        result = func(*args, **kwargs)
+        end = time.perf_counter()
+        elapsed = end - start
+        logger.debug(f"[{func.__name__}] executed in {elapsed:.6f} seconds")
+        return result
+    return wrapper

{mem0ai_azure_mysql-0.1.116.3 → mem0ai_azure_mysql-0.1.116.4}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
 
 [project]
 name = "mem0ai-azure-mysql"
-version = "0.1.116.3"
+version = "0.1.116.4"
 description = "Long-term memory for AI Agents with Azure DefaultAzureCredential authentication and MySQL history database support"
 readme = "README.md"
 requires-python = ">=3.9,<4.0"