mem-llm 2.0.0__py3-none-any.whl

mem_llm/tool_system.py

@@ -0,0 +1,429 @@
+"""
+Tool System for Function Calling
+=================================
+
+Enables agents to call external functions/tools to perform actions.
+Inspired by OpenAI's function calling and LangChain's tool system.
+
+Features:
+- Decorator-based tool definition
+- Automatic schema generation from type hints
+- Tool execution with error handling
+- Tool result formatting
+- Built-in common tools
+
+Author: C. Emre Karataş
+Version: 2.0.0
+"""
+
+import json
+import inspect
+import re
+from typing import Callable, Dict, List, Any, Optional, get_type_hints
+from dataclasses import dataclass, field
+from enum import Enum
+import logging
+
+logger = logging.getLogger(__name__)
+
+
+class ToolCallStatus(Enum):
+    """Status of tool call execution"""
+    SUCCESS = "success"
+    ERROR = "error"
+    NOT_FOUND = "not_found"
+    INVALID_ARGS = "invalid_args"
+
+
+@dataclass
+class ToolParameter:
+    """Tool parameter definition"""
+    name: str
+    type: str
+    description: str
+    required: bool = True
+    default: Any = None
+
+
+@dataclass
+class Tool:
+    """Tool definition"""
+    name: str
+    description: str
+    function: Callable
+    parameters: List[ToolParameter] = field(default_factory=list)
+    return_type: str = "string"
+    category: str = "general"
+    
+    def to_dict(self) -> Dict[str, Any]:
+        """Convert tool to dictionary format for LLM"""
+        return {
+            "name": self.name,
+            "description": self.description,
+            "category": self.category,
+            "parameters": {
+                "type": "object",
+                "properties": {
+                    param.name: {
+                        "type": param.type,
+                        "description": param.description
+                    }
+                    for param in self.parameters
+                },
+                "required": [p.name for p in self.parameters if p.required]
+            },
+            "return_type": self.return_type
+        }
+    
+    def execute(self, **kwargs) -> Any:
+        """Execute the tool with given arguments"""
+        try:
+            # Validate required parameters
+            required_params = [p.name for p in self.parameters if p.required]
+            missing = [p for p in required_params if p not in kwargs]
+            if missing:
+                raise ValueError(f"Missing required parameters: {missing}")
+            
+            # Execute function
+            result = self.function(**kwargs)
+            return result
+        except Exception as e:
+            logger.error(f"Tool execution error ({self.name}): {e}")
+            raise
+
+
+@dataclass
+class ToolCallResult:
+    """Result of a tool call"""
+    tool_name: str
+    status: ToolCallStatus
+    result: Any = None
+    error: Optional[str] = None
+    execution_time: float = 0.0
+
+
+def tool(
+    name: Optional[str] = None,
+    description: Optional[str] = None,
+    category: str = "general"
+) -> Callable:
+    """
+    Decorator to define a tool/function that the agent can call.
+    
+    Args:
+        name: Tool name (defaults to function name)
+        description: Tool description (defaults to docstring)
+        category: Tool category for organization
+    
+    Example:
+        @tool(name="calculate", description="Perform mathematical calculations")
+        def calculator(expression: str) -> float:
+            '''Evaluate a mathematical expression'''
+            return eval(expression)
+    """
+    def decorator(func: Callable) -> Tool:
+        # Get function metadata
+        func_name = name or func.__name__
+        func_desc = description or (func.__doc__ or "").strip()
+        
+        # Extract parameters from type hints
+        type_hints = get_type_hints(func)
+        sig = inspect.signature(func)
+        parameters = []
+        
+        for param_name, param in sig.parameters.items():
+            if param_name in type_hints:
+                param_type = type_hints[param_name]
+                # Map Python types to JSON schema types
+                type_map = {
+                    str: "string",
+                    int: "integer",
+                    float: "number",
+                    bool: "boolean",
+                    list: "array",
+                    dict: "object"
+                }
+                json_type = type_map.get(param_type, "string")
+            else:
+                json_type = "string"
+            
+            param_desc = f"Parameter: {param_name}"
+            required = param.default == inspect.Parameter.empty
+            
+            parameters.append(ToolParameter(
+                name=param_name,
+                type=json_type,
+                description=param_desc,
+                required=required,
+                default=param.default if param.default != inspect.Parameter.empty else None
+            ))
+        
+        # Get return type
+        return_type = "string"
+        if "return" in type_hints:
+            ret_type = type_hints["return"]
+            type_map = {str: "string", int: "integer", float: "number", bool: "boolean"}
+            return_type = type_map.get(ret_type, "string")
+        
+        # Create Tool object
+        tool_obj = Tool(
+            name=func_name,
+            description=func_desc,
+            function=func,
+            parameters=parameters,
+            return_type=return_type,
+            category=category
+        )
+        
+        # Attach tool metadata to function
+        func._tool = tool_obj
+        return func
+    
+    return decorator
+
+
+class ToolRegistry:
+    """Registry for managing available tools"""
+    
+    def __init__(self):
+        self.tools: Dict[str, Tool] = {}
+        self._load_builtin_tools()
+    
+    def _load_builtin_tools(self):
+        """Load built-in tools"""
+        # Import built-in tools when available
+        try:
+            from .builtin_tools import BUILTIN_TOOLS
+            for tool_func in BUILTIN_TOOLS:
+                if hasattr(tool_func, '_tool'):
+                    self.register(tool_func._tool)
+        except ImportError:
+            pass
+    
+    def register(self, tool: Tool):
+        """Register a tool"""
+        self.tools[tool.name] = tool
+        logger.info(f"Registered tool: {tool.name}")
+    
+    def register_function(self, func: Callable):
+        """Register a function as a tool"""
+        if hasattr(func, '_tool'):
+            self.register(func._tool)
+        else:
+            # Auto-create tool from function
+            tool_obj = tool()(func)
+            if hasattr(tool_obj, '_tool'):
+                self.register(tool_obj._tool)
+    
+    def get(self, name: str) -> Optional[Tool]:
+        """Get a tool by name"""
+        return self.tools.get(name)
+    
+    def list_tools(self, category: Optional[str] = None) -> List[Tool]:
+        """List all tools, optionally filtered by category"""
+        tools = list(self.tools.values())
+        if category:
+            tools = [t for t in tools if t.category == category]
+        return tools
+    
+    def get_tools_schema(self) -> List[Dict[str, Any]]:
+        """Get schema for all tools (for LLM prompt)"""
+        return [tool.to_dict() for tool in self.tools.values()]
+    
+    def execute(self, tool_name: str, **kwargs) -> ToolCallResult:
+        """Execute a tool by name"""
+        import time
+        start_time = time.time()
+        
+        # Get tool
+        tool = self.get(tool_name)
+        if not tool:
+            return ToolCallResult(
+                tool_name=tool_name,
+                status=ToolCallStatus.NOT_FOUND,
+                error=f"Tool '{tool_name}' not found",
+                execution_time=time.time() - start_time
+            )
+        
+        # Execute tool
+        try:
+            result = tool.execute(**kwargs)
+            return ToolCallResult(
+                tool_name=tool_name,
+                status=ToolCallStatus.SUCCESS,
+                result=result,
+                execution_time=time.time() - start_time
+            )
+        except ValueError as e:
+            return ToolCallResult(
+                tool_name=tool_name,
+                status=ToolCallStatus.INVALID_ARGS,
+                error=str(e),
+                execution_time=time.time() - start_time
+            )
+        except Exception as e:
+            return ToolCallResult(
+                tool_name=tool_name,
+                status=ToolCallStatus.ERROR,
+                error=str(e),
+                execution_time=time.time() - start_time
+            )
+
+
+class ToolCallParser:
+    """Parse LLM output to detect and extract tool calls"""
+    
+    # Pattern to detect tool calls in LLM output
+    # Format: TOOL_CALL: tool_name(arg1="value1", arg2="value2")
+    TOOL_CALL_PATTERN = r'TOOL_CALL:\s*(\w+)\((.*?)\)'
+    
+    @staticmethod
+    def parse(text: str) -> List[Dict[str, Any]]:
+        """
+        Parse text to extract tool calls.
+        
+        Returns:
+            List of dicts with 'tool' and 'arguments' keys
+        """
+        tool_calls = []
+        
+        # Find all tool call matches
+        matches = re.finditer(ToolCallParser.TOOL_CALL_PATTERN, text, re.MULTILINE)
+        
+        for match in matches:
+            tool_name = match.group(1)
+            args_str = match.group(2)
+            
+            # Parse arguments
+            arguments = {}
+            if args_str.strip():
+                try:
+                    # Try to parse as Python kwargs
+                    # Handle both key="value" and positional args
+                    args_dict = {}
+                    
+                    # Split by comma, but respect quotes and parentheses
+                    parts = []
+                    current = ""
+                    in_quotes = False
+                    paren_depth = 0
+                    quote_char = None
+                    
+                    for char in args_str:
+                        if char in ['"', "'"] and quote_char is None:
+                            quote_char = char
+                            in_quotes = True
+                            current += char
+                        elif char == quote_char:
+                            in_quotes = False
+                            quote_char = None
+                            current += char
+                        elif char == '(' and not in_quotes:
+                            paren_depth += 1
+                            current += char
+                        elif char == ')' and not in_quotes:
+                            paren_depth -= 1
+                            current += char
+                        elif char == ',' and not in_quotes and paren_depth == 0:
+                            if current.strip():
+                                parts.append(current.strip())
+                            current = ""
+                        else:
+                            current += char
+                    
+                    if current.strip():
+                        parts.append(current.strip())
+                    
+                    # Parse each part
+                    for i, part in enumerate(parts):
+                        if '=' in part and not part.strip().startswith('"'):
+                            key, value = part.split('=', 1)
+                            key = key.strip()
+                            value = value.strip().strip('"\'')
+                            args_dict[key] = value
+                        else:
+                            # Positional argument - use index as key
+                            value = part.strip().strip('"\'')
+                            # Try to infer parameter name from common patterns
+                            if i == 0 and value:
+                                # First arg is usually the main parameter
+                                if tool_name == 'calculate':
+                                    args_dict['expression'] = value
+                                elif tool_name in ['count_words', 'reverse_text', 'to_uppercase', 'to_lowercase']:
+                                    args_dict['text'] = value
+                                elif tool_name == 'get_weather':
+                                    args_dict['city'] = value
+                                elif tool_name in ['read_file', 'write_file']:
+                                    args_dict['filepath'] = value
+                                else:
+                                    args_dict[f'arg{i}'] = value
+                    
+                    arguments = args_dict
+                except Exception as e:
+                    logger.warning(f"Failed to parse arguments: {args_str} - Error: {e}")
+            
+            tool_calls.append({
+                "tool": tool_name,
+                "arguments": arguments
+            })
+        
+        return tool_calls
+    
+    @staticmethod
+    def has_tool_call(text: str) -> bool:
+        """Check if text contains a tool call"""
+        return bool(re.search(ToolCallParser.TOOL_CALL_PATTERN, text))
+    
+    @staticmethod
+    def remove_tool_calls(text: str) -> str:
+        """Remove tool call syntax from text, keeping only natural language"""
+        return re.sub(ToolCallParser.TOOL_CALL_PATTERN, '', text).strip()
+
+
+def format_tools_for_prompt(tools: List[Tool]) -> str:
+    """
+    Format tools as a string for LLM prompt.
+    
+    Returns:
+        Formatted string describing available tools
+    """
+    if not tools:
+        return ""
+    
+    lines = ["You have access to the following tools:"]
+    lines.append("")
+    
+    for tool in tools:
+        lines.append(f"- **{tool.name}**: {tool.description}")
+        
+        if tool.parameters:
+            lines.append("  Parameters:")
+            for param in tool.parameters:
+                req = "required" if param.required else "optional"
+                lines.append(f"    - {param.name} ({param.type}, {req}): {param.description}")
+        
+        lines.append("")
+    
+    lines.append("=" * 80)
+    lines.append("TOOL USAGE INSTRUCTIONS:")
+    lines.append("-" * 80)
+    lines.append("To call a tool, use EXACTLY this format:")
+    lines.append('  TOOL_CALL: tool_name(param1="value1", param2="value2")')
+    lines.append("")
+    lines.append("EXAMPLES:")
+    lines.append('  TOOL_CALL: calculate(expression="(25 * 4) + 10")')
+    lines.append('  TOOL_CALL: count_words(text="Hello world from AI")')
+    lines.append('  TOOL_CALL: get_current_time()')
+    lines.append('  TOOL_CALL: read_file(filepath="data.txt")')
+    lines.append("")
+    lines.append("IMPORTANT RULES:")
+    lines.append("  1. Always use named parameters (param=\"value\")")
+    lines.append("  2. Put ALL parameters inside the parentheses")
+    lines.append("  3. Use double quotes for string values")
+    lines.append("  4. One tool call per line")
+    lines.append("  5. After tool execution, you will receive results to continue your response")
+    lines.append("=" * 80)
+    lines.append("")
+    
+    return "\n".join(lines)
+

mem_llm/vector_store.py

@@ -0,0 +1,278 @@
+"""
+Vector Store Abstraction Layer
+Supports multiple vector databases (Chroma, FAISS, etc.)
+"""
+
+from abc import ABC, abstractmethod
+from typing import List, Dict, Optional, Any
+import logging
+
+logger = logging.getLogger(__name__)
+
+
+class VectorStore(ABC):
+    """Abstract interface for vector stores"""
+    
+    @abstractmethod
+    def add_documents(self, documents: List[Dict[str, Any]]) -> None:
+        """
+        Add documents to vector store
+        
+        Args:
+            documents: List of dicts with 'id', 'text', 'metadata'
+        """
+        pass
+    
+    @abstractmethod
+    def search(self, query: str, limit: int = 5, filter_metadata: Optional[Dict] = None) -> List[Dict[str, Any]]:
+        """
+        Search similar documents
+        
+        Args:
+            query: Search query text
+            limit: Maximum number of results
+            filter_metadata: Optional metadata filters
+            
+        Returns:
+            List of similar documents with scores
+        """
+        pass
+    
+    @abstractmethod
+    def delete_collection(self) -> None:
+        """Delete all vectors in collection"""
+        pass
+    
+    @abstractmethod
+    def get_stats(self) -> Dict[str, Any]:
+        """Get statistics about the vector store"""
+        pass
+
+
+try:
+    import chromadb
+    from chromadb.config import Settings
+    CHROMA_AVAILABLE = True
+except ImportError:
+    CHROMA_AVAILABLE = False
+    # Don't warn on import, only when actually trying to use it
+
+
+class ChromaVectorStore(VectorStore):
+    """ChromaDB implementation of VectorStore"""
+    
+    def __init__(self, collection_name: str = "knowledge_base", 
+                 persist_directory: Optional[str] = None,
+                 embedding_model: str = "all-MiniLM-L6-v2"):
+        """
+        Initialize ChromaDB vector store
+        
+        Args:
+            collection_name: Name of the collection
+            persist_directory: Directory to persist data (None = in-memory)
+            embedding_model: Embedding model name (sentence-transformers compatible)
+        """
+        if not CHROMA_AVAILABLE:
+            raise ImportError(
+                "ChromaDB is not installed. Install with: pip install chromadb"
+            )
+        
+        self.collection_name = collection_name
+        self.persist_directory = persist_directory
+        self.embedding_model = embedding_model
+        
+        # Initialize Chroma client
+        if persist_directory:
+            self.client = chromadb.PersistentClient(path=persist_directory)
+        else:
+            self.client = chromadb.Client()
+        
+        # Lazy load embedding model
+        self._embedding_fn = None
+        
+        # Get or create collection with embedding function
+        try:
+            # Create embedding function
+            embedding_fn = self._get_embedding_function()
+            
+            self.collection = self.client.get_or_create_collection(
+                name=collection_name,
+                embedding_function=embedding_fn,
+                metadata={"hnsw:space": "cosine"}
+            )
+        except Exception as e:
+            logger.error(f"Failed to create Chroma collection: {e}")
+            raise
+    
+    def _get_embedding_function(self):
+        """Lazy load embedding function"""
+        if self._embedding_fn is None:
+            try:
+                # Try to use ChromaDB's native SentenceTransformerEmbeddingFunction
+                try:
+                    # Try different import paths for ChromaDB embedding functions
+                    try:
+                        from chromadb.utils import embedding_functions
+                        embedding_fn_class = embedding_functions.SentenceTransformerEmbeddingFunction
+                    except (ImportError, AttributeError):
+                        try:
+                            from chromadb.utils.embedding_functions import SentenceTransformerEmbeddingFunction as embedding_fn_class
+                        except ImportError:
+                            embedding_fn_class = None
+                    
+                    if embedding_fn_class:
+                        self._embedding_fn = embedding_fn_class(model_name=self.embedding_model)
+                        logger.info(f"Loaded embedding model using ChromaDB native function: {self.embedding_model}")
+                    else:
+                        raise AttributeError("SentenceTransformerEmbeddingFunction not found")
+                        
+                except (ImportError, AttributeError, Exception) as e:
+                    # Fallback: Custom embedding function wrapper compatible with ChromaDB
+                    from sentence_transformers import SentenceTransformer
+                    model = SentenceTransformer(self.embedding_model)
+                    
+                    class CustomEmbeddingFunction:
+                        def __init__(self, model, model_name):
+                            self.model = model
+                            self.model_name = model_name
+                            self.name = model_name  # ChromaDB may check for 'name' attribute
+                        
+                        def __call__(self, texts: List[str]) -> List[List[float]]:
+                            embeddings = self.model.encode(texts, show_progress_bar=False)
+                            return embeddings.tolist()
+                        
+                        def encode_queries(self, queries: List[str]) -> List[List[float]]:
+                            return self.__call__(queries)
+                    
+                    self._embedding_fn = CustomEmbeddingFunction(model, self.embedding_model)
+                    logger.info(f"Loaded embedding model using custom wrapper: {self.embedding_model} (fallback: {e})")
+            except ImportError:
+                raise ImportError(
+                    "sentence-transformers not installed. "
+                    "Install with: pip install sentence-transformers"
+                )
+        
+        return self._embedding_fn
+    
+    def add_documents(self, documents: List[Dict[str, Any]]) -> None:
+        """Add documents to ChromaDB"""
+        if not documents:
+            return
+        
+        # Prepare data
+        ids = []
+        texts = []
+        metadatas = []
+        
+        for doc in documents:
+            doc_id = str(doc.get('id', doc.get('text', ''))[:100])
+            # Ensure unique IDs
+            if doc_id in ids:
+                doc_id = f"{doc_id}_{len(ids)}"
+            ids.append(doc_id)
+            texts.append(doc['text'])
+            # Ensure metadata values are JSON-serializable
+            metadata = doc.get('metadata', {})
+            clean_metadata = {}
+            for k, v in metadata.items():
+                if isinstance(v, (str, int, float, bool)) or v is None:
+                    clean_metadata[k] = v
+                else:
+                    clean_metadata[k] = str(v)
+            metadatas.append(clean_metadata)
+        
+        # Add to collection (Chroma will use embedding function automatically)
+        try:
+            self.collection.add(
+                ids=ids,
+                documents=texts,
+                metadatas=metadatas
+            )
+            
+            logger.debug(f"Added {len(documents)} documents to Chroma")
+        except Exception as e:
+            logger.error(f"Error adding documents to Chroma: {e}")
+            raise
+    
+    def search(self, query: str, limit: int = 5, 
+               filter_metadata: Optional[Dict] = None) -> List[Dict[str, Any]]:
+        """Search in ChromaDB"""
+        try:
+            # Build where clause for metadata filtering
+            where = None
+            if filter_metadata:
+                where = filter_metadata
+            
+            # Search (Chroma will use embedding function automatically)
+            results = self.collection.query(
+                query_texts=[query],
+                n_results=limit,
+                where=where
+            )
+            
+            # Format results
+            formatted_results = []
+            if results.get('documents') and len(results['documents']) > 0 and len(results['documents'][0]) > 0:
+                num_results = len(results['documents'][0])
+                distances = results.get('distances', [[0.0] * num_results])
+                
+                for i in range(num_results):
+                    # ChromaDB uses cosine distance (0 = identical, 1 = opposite)
+                    # Convert to similarity score (1 = identical, 0 = opposite)
+                    distance = distances[0][i] if distances and len(distances[0]) > i else 0.0
+                    similarity = 1.0 - distance if distance <= 1.0 else max(0.0, 1.0 / (1.0 + distance))
+                    
+                    formatted_results.append({
+                        'id': results['ids'][0][i] if results.get('ids') and len(results['ids'][0]) > i else f"doc_{i}",
+                        'text': results['documents'][0][i],
+                        'metadata': results['metadatas'][0][i] if results.get('metadatas') and len(results['metadatas'][0]) > i else {},
+                        'score': similarity
+                    })
+            
+            return formatted_results
+        except Exception as e:
+            logger.error(f"Error searching Chroma: {e}")
+            return []
+    
+    def delete_collection(self) -> None:
+        """Delete collection"""
+        try:
+            self.client.delete_collection(self.collection_name)
+            logger.info(f"Deleted Chroma collection: {self.collection_name}")
+        except Exception as e:
+            logger.error(f"Error deleting collection: {e}")
+    
+    def get_stats(self) -> Dict[str, Any]:
+        """Get collection statistics"""
+        try:
+            count = self.collection.count()
+            return {
+                'total_documents': count,
+                'collection_name': self.collection_name,
+                'embedding_model': self.embedding_model
+            }
+        except Exception as e:
+            logger.error(f"Error getting stats: {e}")
+            return {'total_documents': 0}
+
+
+def create_vector_store(store_type: str = "chroma", **kwargs) -> Optional[VectorStore]:
+    """
+    Factory function to create vector store
+    
+    Args:
+        store_type: Type of vector store ('chroma', 'faiss', etc.)
+        **kwargs: Store-specific parameters
+        
+    Returns:
+        VectorStore instance or None if not available
+    """
+    if store_type == "chroma":
+        if not CHROMA_AVAILABLE:
+            logger.info("ℹ️  Vector search disabled (ChromaDB not installed). For semantic search, run: pip install chromadb sentence-transformers")
+            return None
+        return ChromaVectorStore(**kwargs)
+    else:
+        logger.warning(f"Unknown vector store type: {store_type}")
+        return None
+