mem0ai-azure-mysql 0.1.116.1__tar.gz → 0.1.116.2__tar.gz

{mem0ai_azure_mysql-0.1.116.1 → mem0ai_azure_mysql-0.1.116.2}/PKG-INFO

@@ -1,10 +1,11 @@
 Metadata-Version: 2.4
 Name: mem0ai-azure-mysql
-Version: 0.1.116.1
+Version: 0.1.116.2
 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
 Requires-Dist: azure-search-documents>=11.5.3
+Requires-Dist: cohere>=5.17.0
 Requires-Dist: langchain-neo4j>=0.4.0
 Requires-Dist: openai<1.100.0,>=1.90.0
 Requires-Dist: posthog>=3.5.0

{mem0ai_azure_mysql-0.1.116.1 → mem0ai_azure_mysql-0.1.116.2}/mem0/graphs/configs.py

@@ -5,12 +5,24 @@ from pydantic import BaseModel, Field, field_validator, model_validator
 from mem0.llms.configs import LlmConfig
 
 
+class RerankConfig(BaseModel):
+    provider: str = Field(
+        description="Provider of the rerank model (e.g., 'openai', 'azure', 'cohere')",
+        default="cohere", 
+    )
+    config: Optional[dict] = Field(
+        description="Configuration for the specific rerank model", default={}
+    )
+
 class Neo4jConfig(BaseModel):
     url: Optional[str] = Field(None, description="Host address for the graph database")
     username: Optional[str] = Field(None, description="Username for the graph database")
     password: Optional[str] = Field(None, description="Password for the graph database")
     database: Optional[str] = Field(None, description="Database for the graph database")
     base_label: Optional[bool] = Field(None, description="Whether to use base node label __Entity__ for all entities")
+    similarity_threshold: float = Field(0.7, description="Threshold for the similarity of nodes")
+    top_k: int = Field(5, description="Number of top scored results to return")
+    rerank: Optional[RerankConfig] = Field(None, description="Rerank configuration")
 
     @model_validator(mode="before")
     def check_host_port_or_path(cls, values):
@@ -21,6 +33,12 @@ class Neo4jConfig(BaseModel):
         )
         if not url or not username or not password:
             raise ValueError("Please provide 'url', 'username' and 'password'.")
+        
+        if values.get("rerank") is not None:
+            values["rerank"] = RerankConfig(**values.get("rerank"))
+            if values["rerank"].provider not in ("cohere"):
+                raise ValueError("Invalid rerank provider. Supported providers are: cohere")
+
         return values
 
 

{mem0ai_azure_mysql-0.1.116.1 → mem0ai_azure_mysql-0.1.116.2}/mem0/memory/graph_memory.py

@@ -1,4 +1,7 @@
 import logging
+import time
+
+import cohere
 
 from mem0.memory.utils import format_entities, sanitize_relationship_for_cypher
 
@@ -70,7 +73,8 @@ class MemoryGraph:
             llm_config = self.config.llm.config
         self.llm = LlmFactory.create(self.llm_provider, llm_config)
         self.user_id = None
-        self.threshold = 0.7
+        self.threshold = self.config.graph_store.config.similarity_threshold
+        self.top_k = self.config.graph_store.config.top_k
 
     def add(self, data, filters):
         """
@@ -106,28 +110,70 @@ class MemoryGraph:
                 - "contexts": List of search results from the base data store.
                 - "entities": List of related graph data based on the query.
         """
+        start_time = time.perf_counter()
+        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)
 
         if not search_output:
             return []
 
-        search_outputs_sequence = [
-            [item["source"], item["relationship"], item["destination"]] for item in search_output
-        ]
-        bm25 = BM25Okapi(search_outputs_sequence)
-
-        tokenized_query = query.split(" ")
-        reranked_results = bm25.get_top_n(tokenized_query, search_outputs_sequence, n=5)
+        reranked_results = self.rerank(search_output, query)
 
         search_results = []
         for item in reranked_results:
-            search_results.append({"source": item[0], "relationship": item[1], "destination": item[2]})
+            search_results.append(
+                {
+                    "source": item[0],
+                    "relationship": item[1],
+                    "destination": item[2],
+                    "score": item[3] if len(item) > 3 else 0.0,
+                }
+            )
 
         logger.info(f"Returned {len(search_results)} search results")
 
+        end_time = time.perf_counter()
+        # print milliseconds
+        print(f"Search took {(end_time - start_time) * 1000:.2f} ms")
+
         return search_results
 
+    def rerank(self, search_output, query):
+        if not self.config.graph_store.config.rerank:
+            search_outputs_sequence = [
+                [item["source"], item["relationship"], item["destination"]] for item in search_output
+            ]
+            bm25 = BM25Okapi(search_outputs_sequence)
+
+            tokenized_query = query.split(" ")
+            return bm25.get_top_n(tokenized_query, search_outputs_sequence, n=5)
+        
+        # Only cohere is supported for now
+        rerank_config = self.config.graph_store.config.rerank.config
+        co = cohere.ClientV2(
+            api_key=rerank_config.get("api_key"),
+            base_url=rerank_config.get("base_url"),
+        )
+        docs = [f'{item["source"]} {item["relationship"]} {item["destination"]}' for item in search_output]
+        response = co.rerank(
+            model=rerank_config.get("model"),
+            query=query,
+            documents=docs,
+            top_n=self.top_k,
+        )
+
+        reranked_results = []
+        for result in response.results:
+            original_output = search_output[result.index]
+            reranked_results.append([
+                original_output["source"],
+                original_output["relationship"],
+                original_output["destination"],
+                result.relevance_score
+            ])
+        return reranked_results
+
     def delete_all(self, filters):
         # Build node properties for filtering
         node_props = ["user_id: $user_id"]

{mem0ai_azure_mysql-0.1.116.1 → mem0ai_azure_mysql-0.1.116.2}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
 
 [project]
 name = "mem0ai-azure-mysql"
-version = "0.1.116.1"
+version = "0.1.116.2"
 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"
@@ -20,7 +20,8 @@ dependencies = [
     "azure-search-documents>=11.5.3",
     "pymysql",
     "langchain-neo4j>=0.4.0",
-    "rank-bm25>=0.2.2"
+    "rank-bm25>=0.2.2",
+    "cohere>=5.17.0",
 ]
 
 [project.optional-dependencies]