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

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: mem0ai-azure-mysql
-Version: 0.1.116.3
+Version: 0.1.116.5
 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.5}/mem0/embeddings/azure_openai.py

@@ -6,9 +6,13 @@ 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):
+
+    cached_embeddings = {}
+
     def __init__(self, config: Optional[BaseEmbedderConfig] = None):
         super().__init__(config)
 
@@ -41,6 +45,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.
@@ -52,4 +57,28 @@ 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):
+        """
+        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]
+
+        # 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.3 → mem0ai_azure_mysql-0.1.116.5}/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.5}/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,54 +315,54 @@ 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 = []
-
         # 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 node in node_list:
-            n_embedding = self.embedding_model.embed(node)
-
-            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()
@@ -377,7 +374,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 +412,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,12 +468,46 @@ 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"]
         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)}
+
+        query_stats = {
+                1: 0, 2: 0, 3: 0, 4: 0
+            }
+    
+        queries = {
+            "only_source_found": {
+                "template": "",
+                "rows": []
+            },
+            "only_destination_found": {
+                "template": "",
+                "rows": []
+            },
+            "both_found": {
+                "template": "",
+                "rows": []
+            },
+            "none_found": {
+                "template": "",
+                "rows": []
+            },
+        }
+
         for item in to_be_added:
             # entities
             source = item["source"]
@@ -490,12 +523,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:
@@ -541,6 +574,8 @@ class MemoryGraph:
                     params["agent_id"] = agent_id
                 if run_id:
                     params["run_id"] = run_id
+                
+                query_stats[1] += 1
 
             elif destination_node_search_result and not source_node_search_result:
                 # Build source MERGE properties
@@ -585,6 +620,8 @@ class MemoryGraph:
                     params["agent_id"] = agent_id
                 if run_id:
                     params["run_id"] = run_id
+                
+                query_stats[2] += 1
 
             elif source_node_search_result and destination_node_search_result:
                 cypher = f"""
@@ -613,6 +650,8 @@ class MemoryGraph:
                     params["agent_id"] = agent_id
                 if run_id:
                     params["run_id"] = run_id
+                
+                query_stats[3] += 1
 
             else:
                 # Build dynamic MERGE props for both source and destination
@@ -661,8 +700,12 @@ class MemoryGraph:
                     params["agent_id"] = agent_id
                 if run_id:
                     params["run_id"] = run_id
+                
+                query_stats[4] += 1
+
             result = self.graph.query(cypher, params=params)
             results.append(result)
+        logger.warning(f"Add entities query stats: {query_stats}")
         return results
 
     def _remove_spaces_from_entities(self, entity_list):
@@ -673,6 +716,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"]
@@ -710,6 +754,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.3 → mem0ai_azure_mysql-0.1.116.5}/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:
@@ -1119,6 +1122,7 @@ class AsyncMemory(MemoryBase):
 
         return {"results": vector_store_result}
 
+    @time_perf
     async def _add_to_vector_store(
         self,
         messages: list,
@@ -1303,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:

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

@@ -1,8 +1,13 @@
-import hashlib
 import re
+import time
+import hashlib
+import logging
+import asyncio
+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 +187,26 @@ def sanitize_relationship_for_cypher(relationship) -> str:
         sanitized = sanitized.replace(old, new)
 
     return re.sub(r"_+", "_", sanitized).strip("_")
+
+
+def time_perf(func):
+    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.warning(f"[{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.warning(f"[{func.__name__}] executed in {elapsed:.6f} seconds")
+            return result
+        return sync_wrapper

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

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