agentrun-mem0ai 0.0.11__py3-none-any.whl

agentrun_mem0/proxy/main.py

@@ -0,0 +1,189 @@
+import logging
+import subprocess
+import sys
+import threading
+from typing import List, Optional, Union
+
+import httpx
+
+import agentrun_mem0
+
+try:
+    import litellm
+except ImportError:
+    try:
+        subprocess.check_call([sys.executable, "-m", "pip", "install", "litellm"])
+        import litellm
+    except subprocess.CalledProcessError:
+        print("Failed to install 'litellm'. Please install it manually using 'pip install litellm'.")
+        sys.exit(1)
+
+from agentrun_mem0 import Memory, MemoryClient
+from agentrun_mem0.configs.prompts import MEMORY_ANSWER_PROMPT
+from agentrun_mem0.memory.telemetry import capture_client_event, capture_event
+
+logger = logging.getLogger(__name__)
+
+
+class Mem0:
+    def __init__(
+        self,
+        config: Optional[dict] = None,
+        api_key: Optional[str] = None,
+        host: Optional[str] = None,
+    ):
+        if api_key:
+            self.mem0_client = MemoryClient(api_key, host)
+        else:
+            self.mem0_client = Memory.from_config(config) if config else Memory()
+
+        self.chat = Chat(self.mem0_client)
+
+
+class Chat:
+    def __init__(self, mem0_client):
+        self.completions = Completions(mem0_client)
+
+
+class Completions:
+    def __init__(self, mem0_client):
+        self.mem0_client = mem0_client
+
+    def create(
+        self,
+        model: str,
+        messages: List = [],
+        # Mem0 arguments
+        user_id: Optional[str] = None,
+        agent_id: Optional[str] = None,
+        run_id: Optional[str] = None,
+        metadata: Optional[dict] = None,
+        filters: Optional[dict] = None,
+        limit: Optional[int] = 10,
+        # LLM arguments
+        timeout: Optional[Union[float, str, httpx.Timeout]] = None,
+        temperature: Optional[float] = None,
+        top_p: Optional[float] = None,
+        n: Optional[int] = None,
+        stream: Optional[bool] = None,
+        stream_options: Optional[dict] = None,
+        stop=None,
+        max_tokens: Optional[int] = None,
+        presence_penalty: Optional[float] = None,
+        frequency_penalty: Optional[float] = None,
+        logit_bias: Optional[dict] = None,
+        user: Optional[str] = None,
+        # openai v1.0+ new params
+        response_format: Optional[dict] = None,
+        seed: Optional[int] = None,
+        tools: Optional[List] = None,
+        tool_choice: Optional[Union[str, dict]] = None,
+        logprobs: Optional[bool] = None,
+        top_logprobs: Optional[int] = None,
+        parallel_tool_calls: Optional[bool] = None,
+        deployment_id=None,
+        extra_headers: Optional[dict] = None,
+        # soon to be deprecated params by OpenAI
+        functions: Optional[List] = None,
+        function_call: Optional[str] = None,
+        # set api_base, api_version, api_key
+        base_url: Optional[str] = None,
+        api_version: Optional[str] = None,
+        api_key: Optional[str] = None,
+        model_list: Optional[list] = None,  # pass in a list of api_base,keys, etc.
+    ):
+        if not any([user_id, agent_id, run_id]):
+            raise ValueError("One of user_id, agent_id, run_id must be provided")
+
+        if not litellm.supports_function_calling(model):
+            raise ValueError(
+                f"Model '{model}' does not support function calling. Please use a model that supports function calling."
+            )
+
+        prepared_messages = self._prepare_messages(messages)
+        if prepared_messages[-1]["role"] == "user":
+            self._async_add_to_memory(messages, user_id, agent_id, run_id, metadata, filters)
+            relevant_memories = self._fetch_relevant_memories(messages, user_id, agent_id, run_id, filters, limit)
+            logger.debug(f"Retrieved {len(relevant_memories)} relevant memories")
+            prepared_messages[-1]["content"] = self._format_query_with_memories(messages, relevant_memories)
+
+        response = litellm.completion(
+            model=model,
+            messages=prepared_messages,
+            temperature=temperature,
+            top_p=top_p,
+            n=n,
+            timeout=timeout,
+            stream=stream,
+            stream_options=stream_options,
+            stop=stop,
+            max_tokens=max_tokens,
+            presence_penalty=presence_penalty,
+            frequency_penalty=frequency_penalty,
+            logit_bias=logit_bias,
+            user=user,
+            response_format=response_format,
+            seed=seed,
+            tools=tools,
+            tool_choice=tool_choice,
+            logprobs=logprobs,
+            top_logprobs=top_logprobs,
+            parallel_tool_calls=parallel_tool_calls,
+            deployment_id=deployment_id,
+            extra_headers=extra_headers,
+            functions=functions,
+            function_call=function_call,
+            base_url=base_url,
+            api_version=api_version,
+            api_key=api_key,
+            model_list=model_list,
+        )
+        if isinstance(self.mem0_client, Memory):
+            capture_event("mem0.chat.create", self.mem0_client)
+        else:
+            capture_client_event("mem0.chat.create", self.mem0_client)
+        return response
+
+    def _prepare_messages(self, messages: List[dict]) -> List[dict]:
+        if not messages or messages[0]["role"] != "system":
+            return [{"role": "system", "content": MEMORY_ANSWER_PROMPT}] + messages
+        return messages
+
+    def _async_add_to_memory(self, messages, user_id, agent_id, run_id, metadata, filters):
+        def add_task():
+            logger.debug("Adding to memory asynchronously")
+            self.mem0_client.add(
+                messages=messages,
+                user_id=user_id,
+                agent_id=agent_id,
+                run_id=run_id,
+                metadata=metadata,
+                filters=filters,
+            )
+
+        threading.Thread(target=add_task, daemon=True).start()
+
+    def _fetch_relevant_memories(self, messages, user_id, agent_id, run_id, filters, limit):
+        # Currently, only pass the last 6 messages to the search API to prevent long query
+        message_input = [f"{message['role']}: {message['content']}" for message in messages][-6:]
+        # TODO: Make it better by summarizing the past conversation
+        return self.mem0_client.search(
+            query="\n".join(message_input),
+            user_id=user_id,
+            agent_id=agent_id,
+            run_id=run_id,
+            filters=filters,
+            limit=limit,
+        )
+
+    def _format_query_with_memories(self, messages, relevant_memories):
+        # Check if self.mem0_client is an instance of Memory or MemoryClient
+
+        entities = []
+        if isinstance(self.mem0_client, agentrun_mem0.memory.main.Memory):
+            memories_text = "\n".join(memory["memory"] for memory in relevant_memories["results"])
+            if relevant_memories.get("relations"):
+                entities = [entity for entity in relevant_memories["relations"]]
+        elif isinstance(self.mem0_client, agentrun_mem0.client.main.MemoryClient):
+            memories_text = "\n".join(memory["memory"] for memory in relevant_memories)
+        return f"- Relevant Memories/Facts: {memories_text}\n\n- Entities: {entities}\n\n- User Question: {messages[-1]['content']}"

agentrun_mem0/reranker/__init__.py

@@ -0,0 +1,9 @@
+"""
+Reranker implementations for mem0 search functionality.
+"""
+
+from .base import BaseReranker
+from .cohere_reranker import CohereReranker
+from .sentence_transformer_reranker import SentenceTransformerReranker
+
+__all__ = ["BaseReranker", "CohereReranker", "SentenceTransformerReranker"]

agentrun_mem0/reranker/base.py

@@ -0,0 +1,20 @@
+from abc import ABC, abstractmethod
+from typing import List, Dict, Any
+
+class BaseReranker(ABC):
+    """Abstract base class for all rerankers."""
+    
+    @abstractmethod
+    def rerank(self, query: str, documents: List[Dict[str, Any]], top_k: int = None) -> List[Dict[str, Any]]:
+        """
+        Rerank documents based on relevance to the query.
+        
+        Args:
+            query: The search query
+            documents: List of documents to rerank, each with 'memory' field  
+            top_k: Number of top documents to return (None = return all)
+            
+        Returns:
+            List of reranked documents with added 'rerank_score' field
+        """
+        pass

agentrun_mem0/reranker/cohere_reranker.py

@@ -0,0 +1,85 @@
+import os
+from typing import List, Dict, Any
+
+from agentrun_mem0.reranker.base import BaseReranker
+
+try:
+    import cohere
+    COHERE_AVAILABLE = True
+except ImportError:
+    COHERE_AVAILABLE = False
+
+
+class CohereReranker(BaseReranker):
+    """Cohere-based reranker implementation."""
+    
+    def __init__(self, config):
+        """
+        Initialize Cohere reranker.
+        
+        Args:
+            config: CohereRerankerConfig object with configuration parameters
+        """
+        if not COHERE_AVAILABLE:
+            raise ImportError("cohere package is required for CohereReranker. Install with: pip install cohere")
+        
+        self.config = config
+        self.api_key = config.api_key or os.getenv("COHERE_API_KEY")
+        if not self.api_key:
+            raise ValueError("Cohere API key is required. Set COHERE_API_KEY environment variable or pass api_key in config.")
+            
+        self.model = config.model
+        self.client = cohere.Client(self.api_key)
+        
+    def rerank(self, query: str, documents: List[Dict[str, Any]], top_k: int = None) -> List[Dict[str, Any]]:
+        """
+        Rerank documents using Cohere's rerank API.
+        
+        Args:
+            query: The search query
+            documents: List of documents to rerank
+            top_k: Number of top documents to return
+            
+        Returns:
+            List of reranked documents with rerank_score
+        """
+        if not documents:
+            return documents
+            
+        # Extract text content for reranking
+        doc_texts = []
+        for doc in documents:
+            if 'memory' in doc:
+                doc_texts.append(doc['memory'])
+            elif 'text' in doc:
+                doc_texts.append(doc['text'])  
+            elif 'content' in doc:
+                doc_texts.append(doc['content'])
+            else:
+                doc_texts.append(str(doc))
+        
+        try:
+            # Call Cohere rerank API
+            response = self.client.rerank(
+                model=self.model,
+                query=query,
+                documents=doc_texts,
+                top_n=top_k or self.config.top_k or len(documents),
+                return_documents=self.config.return_documents,
+                max_chunks_per_doc=self.config.max_chunks_per_doc,
+            )
+            
+            # Create reranked results
+            reranked_docs = []
+            for result in response.results:
+                original_doc = documents[result.index].copy()
+                original_doc['rerank_score'] = result.relevance_score
+                reranked_docs.append(original_doc)
+                
+            return reranked_docs
+
+        except Exception:
+            # Fallback to original order if reranking fails
+            for doc in documents:
+                doc['rerank_score'] = 0.0
+            return documents[:top_k] if top_k else documents

agentrun_mem0/reranker/huggingface_reranker.py

@@ -0,0 +1,147 @@
+from typing import List, Dict, Any, Union
+import numpy as np
+
+from agentrun_mem0.reranker.base import BaseReranker
+from agentrun_mem0.configs.rerankers.base import BaseRerankerConfig
+from agentrun_mem0.configs.rerankers.huggingface import HuggingFaceRerankerConfig
+
+try:
+    from transformers import AutoTokenizer, AutoModelForSequenceClassification
+    import torch
+    TRANSFORMERS_AVAILABLE = True
+except ImportError:
+    TRANSFORMERS_AVAILABLE = False
+
+
+class HuggingFaceReranker(BaseReranker):
+    """HuggingFace Transformers based reranker implementation."""
+
+    def __init__(self, config: Union[BaseRerankerConfig, HuggingFaceRerankerConfig, Dict]):
+        """
+        Initialize HuggingFace reranker.
+
+        Args:
+            config: Configuration object with reranker parameters
+        """
+        if not TRANSFORMERS_AVAILABLE:
+            raise ImportError("transformers package is required for HuggingFaceReranker. Install with: pip install transformers torch")
+
+        # Convert to HuggingFaceRerankerConfig if needed
+        if isinstance(config, dict):
+            config = HuggingFaceRerankerConfig(**config)
+        elif isinstance(config, BaseRerankerConfig) and not isinstance(config, HuggingFaceRerankerConfig):
+            # Convert BaseRerankerConfig to HuggingFaceRerankerConfig with defaults
+            config = HuggingFaceRerankerConfig(
+                provider=getattr(config, 'provider', 'huggingface'),
+                model=getattr(config, 'model', 'BAAI/bge-reranker-base'),
+                api_key=getattr(config, 'api_key', None),
+                top_k=getattr(config, 'top_k', None),
+                device=None,  # Will auto-detect
+                batch_size=32,  # Default
+                max_length=512,  # Default
+                normalize=True,  # Default
+            )
+
+        self.config = config
+
+        # Set device
+        if self.config.device is None:
+            self.device = "cuda" if torch.cuda.is_available() else "cpu"
+        else:
+            self.device = self.config.device
+
+        # Load model and tokenizer
+        self.tokenizer = AutoTokenizer.from_pretrained(self.config.model)
+        self.model = AutoModelForSequenceClassification.from_pretrained(self.config.model)
+        self.model.to(self.device)
+        self.model.eval()
+
+    def rerank(self, query: str, documents: List[Dict[str, Any]], top_k: int = None) -> List[Dict[str, Any]]:
+        """
+        Rerank documents using HuggingFace cross-encoder model.
+
+        Args:
+            query: The search query
+            documents: List of documents to rerank
+            top_k: Number of top documents to return
+
+        Returns:
+            List of reranked documents with rerank_score
+        """
+        if not documents:
+            return documents
+
+        # Extract text content for reranking
+        doc_texts = []
+        for doc in documents:
+            if 'memory' in doc:
+                doc_texts.append(doc['memory'])
+            elif 'text' in doc:
+                doc_texts.append(doc['text'])
+            elif 'content' in doc:
+                doc_texts.append(doc['content'])
+            else:
+                doc_texts.append(str(doc))
+
+        try:
+            scores = []
+
+            # Process documents in batches
+            for i in range(0, len(doc_texts), self.config.batch_size):
+                batch_docs = doc_texts[i:i + self.config.batch_size]
+                batch_pairs = [[query, doc] for doc in batch_docs]
+
+                # Tokenize batch
+                inputs = self.tokenizer(
+                    batch_pairs,
+                    padding=True,
+                    truncation=True,
+                    max_length=self.config.max_length,
+                    return_tensors="pt"
+                ).to(self.device)
+
+                # Get scores
+                with torch.no_grad():
+                    outputs = self.model(**inputs)
+                    batch_scores = outputs.logits.squeeze(-1).cpu().numpy()
+
+                    # Handle single item case
+                    if batch_scores.ndim == 0:
+                        batch_scores = [float(batch_scores)]
+                    else:
+                        batch_scores = batch_scores.tolist()
+
+                    scores.extend(batch_scores)
+
+            # Normalize scores if requested
+            if self.config.normalize:
+                scores = np.array(scores)
+                scores = (scores - scores.min()) / (scores.max() - scores.min() + 1e-8)
+                scores = scores.tolist()
+
+            # Combine documents with scores
+            doc_score_pairs = list(zip(documents, scores))
+
+            # Sort by score (descending)
+            doc_score_pairs.sort(key=lambda x: x[1], reverse=True)
+
+            # Apply top_k limit
+            final_top_k = top_k or self.config.top_k
+            if final_top_k:
+                doc_score_pairs = doc_score_pairs[:final_top_k]
+
+            # Create reranked results
+            reranked_docs = []
+            for doc, score in doc_score_pairs:
+                reranked_doc = doc.copy()
+                reranked_doc['rerank_score'] = float(score)
+                reranked_docs.append(reranked_doc)
+
+            return reranked_docs
+
+        except Exception:
+            # Fallback to original order if reranking fails
+            for doc in documents:
+                doc['rerank_score'] = 0.0
+            final_top_k = top_k or self.config.top_k
+            return documents[:final_top_k] if final_top_k else documents

agentrun_mem0/reranker/llm_reranker.py

@@ -0,0 +1,142 @@
+import re
+from typing import List, Dict, Any, Union
+
+from agentrun_mem0.reranker.base import BaseReranker
+from agentrun_mem0.utils.factory import LlmFactory
+from agentrun_mem0.configs.rerankers.base import BaseRerankerConfig
+from agentrun_mem0.configs.rerankers.llm import LLMRerankerConfig
+
+
+class LLMReranker(BaseReranker):
+    """LLM-based reranker implementation."""
+
+    def __init__(self, config: Union[BaseRerankerConfig, LLMRerankerConfig, Dict]):
+        """
+        Initialize LLM reranker.
+
+        Args:
+            config: Configuration object with reranker parameters
+        """
+        # Convert to LLMRerankerConfig if needed
+        if isinstance(config, dict):
+            config = LLMRerankerConfig(**config)
+        elif isinstance(config, BaseRerankerConfig) and not isinstance(config, LLMRerankerConfig):
+            # Convert BaseRerankerConfig to LLMRerankerConfig with defaults
+            config = LLMRerankerConfig(
+                provider=getattr(config, 'provider', 'openai'),
+                model=getattr(config, 'model', 'gpt-4o-mini'),
+                api_key=getattr(config, 'api_key', None),
+                top_k=getattr(config, 'top_k', None),
+                temperature=0.0,  # Default for reranking
+                max_tokens=100,   # Default for reranking
+            )
+
+        self.config = config
+
+        # Create LLM configuration for the factory
+        llm_config = {
+            "model": self.config.model,
+            "temperature": self.config.temperature,
+            "max_tokens": self.config.max_tokens,
+        }
+
+        # Add API key if provided
+        if self.config.api_key:
+            llm_config["api_key"] = self.config.api_key
+
+        # Initialize LLM using the factory
+        self.llm = LlmFactory.create(self.config.provider, llm_config)
+
+        # Default scoring prompt
+        self.scoring_prompt = getattr(self.config, 'scoring_prompt', None) or self._get_default_prompt()
+        
+    def _get_default_prompt(self) -> str:
+        """Get the default scoring prompt template."""
+        return """You are a relevance scoring assistant. Given a query and a document, you need to score how relevant the document is to the query.
+
+Score the relevance on a scale from 0.0 to 1.0, where:
+- 1.0 = Perfectly relevant and directly answers the query
+- 0.8-0.9 = Highly relevant with good information
+- 0.6-0.7 = Moderately relevant with some useful information  
+- 0.4-0.5 = Slightly relevant with limited useful information
+- 0.0-0.3 = Not relevant or no useful information
+
+Query: "{query}"
+Document: "{document}"
+
+Provide only a single numerical score between 0.0 and 1.0. Do not include any explanation or additional text."""
+
+    def _extract_score(self, response_text: str) -> float:
+        """Extract numerical score from LLM response."""
+        # Look for decimal numbers between 0.0 and 1.0
+        pattern = r'\b([01](?:\.\d+)?)\b'
+        matches = re.findall(pattern, response_text)
+        
+        if matches:
+            score = float(matches[0])
+            return min(max(score, 0.0), 1.0)  # Clamp between 0.0 and 1.0
+        
+        # Fallback: return 0.5 if no valid score found
+        return 0.5
+    
+    def rerank(self, query: str, documents: List[Dict[str, Any]], top_k: int = None) -> List[Dict[str, Any]]:
+        """
+        Rerank documents using LLM scoring.
+        
+        Args:
+            query: The search query
+            documents: List of documents to rerank
+            top_k: Number of top documents to return
+            
+        Returns:
+            List of reranked documents with rerank_score
+        """
+        if not documents:
+            return documents
+        
+        scored_docs = []
+        
+        for doc in documents:
+            # Extract text content
+            if 'memory' in doc:
+                doc_text = doc['memory']
+            elif 'text' in doc:
+                doc_text = doc['text']  
+            elif 'content' in doc:
+                doc_text = doc['content']
+            else:
+                doc_text = str(doc)
+            
+            try:
+                # Generate scoring prompt
+                prompt = self.scoring_prompt.format(query=query, document=doc_text)
+                
+                # Get LLM response
+                response = self.llm.generate_response(
+                    messages=[{"role": "user", "content": prompt}]
+                )
+                
+                # Extract score from response
+                score = self._extract_score(response)
+                
+                # Create scored document
+                scored_doc = doc.copy()
+                scored_doc['rerank_score'] = score
+                scored_docs.append(scored_doc)
+
+            except Exception:
+                # Fallback: assign neutral score if scoring fails
+                scored_doc = doc.copy()
+                scored_doc['rerank_score'] = 0.5
+                scored_docs.append(scored_doc)
+        
+        # Sort by relevance score in descending order
+        scored_docs.sort(key=lambda x: x['rerank_score'], reverse=True)
+        
+        # Apply top_k limit
+        if top_k:
+            scored_docs = scored_docs[:top_k]
+        elif self.config.top_k:
+            scored_docs = scored_docs[:self.config.top_k]
+            
+        return scored_docs