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

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: mem0ai-azure-mysql
-Version: 0.1.116.4
+Version: 0.1.116.6
 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.4 → mem0ai_azure_mysql-0.1.116.6}/mem0/embeddings/azure_openai.py

@@ -10,6 +10,9 @@ from mem0.memory.utils import time_perf
 
 
 class AzureOpenAIEmbedding(EmbeddingBase):
+
+    cached_embeddings = {}
+
     def __init__(self, config: Optional[BaseEmbedderConfig] = None):
         super().__init__(config)
 
@@ -54,7 +57,9 @@ class AzureOpenAIEmbedding(EmbeddingBase):
             list: The embedding vector.
         """
         text = text.replace("\n", " ")
-        return self.client.embeddings.create(input=[text], model=self.config.model).data[0].embedding
+        if text not in self.cached_embeddings:
+            self.cached_embeddings[text] = self.client.embeddings.create(input=[text], model=self.config.model).data[0].embedding
+        return self.cached_embeddings[text]
 
     @time_perf
     def embed_in_batch(self, texts, memory_action: Optional[Literal["add", "search", "update"]] = None):
@@ -68,5 +73,12 @@ class AzureOpenAIEmbedding(EmbeddingBase):
             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]
+
+        # Use cached embeddings if available
+        uncached_texts = [text for text in texts if text not in self.cached_embeddings]
+        if uncached_texts:
+            uncached_embeddings = self.client.embeddings.create(input=uncached_texts, model=self.config.model)
+            for text, item in zip(uncached_texts, uncached_embeddings.data):
+                self.cached_embeddings[text] = item.embedding
+
+        return [self.cached_embeddings[text] for text in texts]

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

@@ -96,6 +96,7 @@ class MemoryGraph:
 
         return {"deleted_entities": deleted_entities, "added_entities": added_entities}
 
+    @time_perf
     def search(self, query, filters, limit=100):
         """
         Search for memories and related graph data.
@@ -139,6 +140,7 @@ class MemoryGraph:
 
         return search_results
 
+    @time_perf
     def rerank(self, search_output, query):
         if not self.config.graph_store.config.rerank:
             search_outputs_sequence = [
@@ -318,52 +320,51 @@ class MemoryGraph:
     @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 = []
-
         # Build node properties for filtering
-        node_props = ["user_id: $user_id"]
+        node_props = ["user_id: row.user_id"]
         if filters.get("agent_id"):
-            node_props.append("agent_id: $agent_id")
+            node_props.append("agent_id: row.agent_id")
         if filters.get("run_id"):
-            node_props.append("run_id: $run_id")
+            node_props.append("run_id: row.run_id")
         node_props_str = ", ".join(node_props)
         n_embeddings = self.embedding_model.embed_in_batch(node_list)
 
-        for n_embedding in n_embeddings:
-
-            cypher_query = f"""
-            MATCH (n {self.node_label} {{{node_props_str}}})
-            WHERE n.embedding IS NOT NULL
-            WITH n, round(2 * vector.similarity.cosine(n.embedding, $n_embedding) - 1, 4) AS similarity // denormalize for backward compatibility
-            WHERE similarity >= $threshold
-            CALL {{
-                WITH n
-                MATCH (n)-[r]->(m {self.node_label} {{{node_props_str}}})
-                RETURN n.name AS source, elementId(n) AS source_id, type(r) AS relationship, elementId(r) AS relation_id, m.name AS destination, elementId(m) AS destination_id
-                UNION
-                WITH n  
-                MATCH (n)<-[r]-(m {self.node_label} {{{node_props_str}}})
-                RETURN m.name AS source, elementId(m) AS source_id, type(r) AS relationship, elementId(r) AS relation_id, n.name AS destination, elementId(n) AS destination_id
-            }}
-            WITH distinct source, source_id, relationship, relation_id, destination, destination_id, similarity
-            RETURN source, source_id, relationship, relation_id, destination, destination_id, similarity
-            ORDER BY similarity DESC
-            LIMIT $limit
-            """
+        cypher_query = f"""
+        UNWIND $rows AS row
+        MATCH (n {self.node_label} {{{node_props_str}}})
+        WHERE n.embedding IS NOT NULL
+        WITH n, row, round(2 * vector.similarity.cosine(n.embedding, row.n_embedding) - 1, 4) AS similarity // denormalize for backward compatibility
+        WHERE similarity >= row.threshold
+        CALL {{
+            WITH n, row
+            MATCH (n)-[r]->(m {self.node_label} {{{node_props_str}}})
+            RETURN n.name AS source, elementId(n) AS source_id, type(r) AS relationship, elementId(r) AS relation_id, m.name AS destination, elementId(m) AS destination_id
+            UNION
+            WITH n, row
+            MATCH (n)<-[r]-(m {self.node_label} {{{node_props_str}}})
+            RETURN m.name AS source, elementId(m) AS source_id, type(r) AS relationship, elementId(r) AS relation_id, n.name AS destination, elementId(n) AS destination_id
+        }}
+        WITH distinct source, source_id, relationship, relation_id, destination, destination_id, similarity
+        RETURN source, source_id, relationship, relation_id, destination, destination_id, similarity
+        ORDER BY similarity DESC
+        LIMIT {limit}
+        """
 
-            params = {
+        rows = []
+        for n_embedding in n_embeddings:
+            row = {
                 "n_embedding": n_embedding,
                 "threshold": self.threshold,
                 "user_id": filters["user_id"],
-                "limit": limit,
             }
             if filters.get("agent_id"):
-                params["agent_id"] = filters["agent_id"]
+                row["agent_id"] = filters["agent_id"]
             if filters.get("run_id"):
-                params["run_id"] = filters["run_id"]
-
-            ans = self.graph.query(cypher_query, params=params)
-            result_relations.extend(ans)
+                row["run_id"] = filters["run_id"]
+            
+            rows.append(row)
+        
+        result_relations = self.graph.query(cypher_query, params={"rows": rows})
 
         # Remove duplicates, key is (source, relationship, destination)
         seen = set()
@@ -476,6 +477,16 @@ class MemoryGraph:
         agent_id = filters.get("agent_id", None)
         run_id = filters.get("run_id", None)
         results = []
+
+        # Collect all entities which needs to be embedded
+        entities_to_embed = set()
+        for item in to_be_added:
+            entities_to_embed.add(item["source"])
+            entities_to_embed.add(item["destination"])
+        entities_to_embed = list(entities_to_embed)
+        embeddings = self.embedding_model.embed_in_batch(entities_to_embed)
+        embedding_map = {entity: embedding for entity, embedding in zip(entities_to_embed, embeddings)}
+
         for item in to_be_added:
             # entities
             source = item["source"]
@@ -491,12 +502,12 @@ class MemoryGraph:
             destination_extra_set = f", destination:`{destination_type}`" if self.node_label else ""
 
             # embeddings
-            source_embedding = self.embedding_model.embed(source)
-            dest_embedding = self.embedding_model.embed(destination)
+            source_embedding = embedding_map[source]
+            dest_embedding = embedding_map[destination]
 
             # search for the nodes with the closest embeddings
-            source_node_search_result = self._search_source_node(source_embedding, filters, threshold=0.9)
-            destination_node_search_result = self._search_destination_node(dest_embedding, filters, threshold=0.9)
+            source_node_search_result = self._search_source_node(source_embedding, filters, threshold=self.threshold)
+            destination_node_search_result = self._search_destination_node(dest_embedding, filters, threshold=self.threshold)
 
             # TODO: Create a cypher query and common params for all the cases
             if not destination_node_search_result and source_node_search_result:
@@ -674,6 +685,7 @@ class MemoryGraph:
             item["destination"] = item["destination"].lower().replace(" ", "_")
         return entity_list
 
+    @time_perf
     def _search_source_node(self, source_embedding, filters, threshold=0.9):
         # Build WHERE conditions
         where_conditions = ["source_candidate.embedding IS NOT NULL", "source_candidate.user_id = $user_id"]
@@ -711,6 +723,7 @@ class MemoryGraph:
         result = self.graph.query(cypher, params=params)
         return result
 
+    @time_perf
     def _search_destination_node(self, destination_embedding, filters, threshold=0.9):
         # Build WHERE conditions
         where_conditions = ["destination_candidate.embedding IS NOT NULL", "destination_candidate.user_id = $user_id"]

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

@@ -1122,6 +1122,7 @@ class AsyncMemory(MemoryBase):
 
         return {"results": vector_store_result}
 
+    @time_perf
     async def _add_to_vector_store(
         self,
         messages: list,
@@ -1306,6 +1307,7 @@ class AsyncMemory(MemoryBase):
         )
         return returned_memories
 
+    @time_perf
     async def _add_to_graph(self, messages, filters):
         added_entities = []
         if self.enable_graph:
@@ -1551,6 +1553,7 @@ class AsyncMemory(MemoryBase):
         else:
             return {"results": original_memories}
 
+    @time_perf
     async def _search_vector_store(self, query, filters, limit, threshold: Optional[float] = None):
         embeddings = await asyncio.to_thread(self.embedding_model.embed, query, "search")
         memories = await asyncio.to_thread(

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

@@ -2,6 +2,7 @@ import re
 import time
 import hashlib
 import logging
+import asyncio
 from functools import wraps
 
 from mem0.configs.prompts import FACT_RETRIEVAL_PROMPT
@@ -189,12 +190,23 @@ def sanitize_relationship_for_cypher(relationship) -> str:
 
 
 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
+    if asyncio.iscoroutinefunction(func):
+        @wraps(func)
+        async def async_wrapper(*args, **kwargs):
+            start = time.perf_counter()
+            result = await func(*args, **kwargs)
+            end = time.perf_counter()
+            elapsed = end - start
+            logger.debug(f"[Mem0 time perf][{func.__name__}] executed in {elapsed:.6f} seconds")
+            return result
+        return async_wrapper
+    else:
+        @wraps(func)
+        def sync_wrapper(*args, **kwargs):
+            start = time.perf_counter()
+            result = func(*args, **kwargs)
+            end = time.perf_counter()
+            elapsed = end - start
+            logger.debug(f"[Mem0 time perf][{func.__name__}] executed in {elapsed:.6f} seconds")
+            return result
+        return sync_wrapper

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

@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
 
 [project]
 name = "mem0ai-azure-mysql"
-version = "0.1.116.4"
+version = "0.1.116.6"
 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"