quantalogic 0.35.0__py3-none-any.whl → 0.40.0__py3-none-any.whl

quantalogic/tools/rag_tool/rag_tool_beta.py

@@ -0,0 +1,264 @@
+"""RAG (Retrieval Augmented Generation) Tool using LlamaIndex.
+
+This tool provides a flexible RAG implementation supporting multiple vector stores
+and embedding models, with configurable document processing options.
+"""
+
+import os
+from enum import Enum
+from typing import Any, List, Optional
+
+import chromadb
+from llama_index.core import (
+    SimpleDirectoryReader,
+    StorageContext,
+    VectorStoreIndex,
+    load_index_from_storage,
+)
+from llama_index.core.settings import Settings
+from llama_index.embeddings.bedrock import BedrockEmbedding
+from llama_index.embeddings.huggingface import HuggingFaceEmbedding
+from llama_index.embeddings.instructor import InstructorEmbedding
+from llama_index.embeddings.openai import OpenAIEmbedding
+from llama_index.vector_stores.chroma import ChromaVectorStore
+from llama_index.vector_stores.faiss import FaissVectorStore
+from loguru import logger
+
+from quantalogic.tools.tool import Tool, ToolArgument
+
+
+class VectorStoreType(str, Enum):
+    """Supported vector store types."""
+    CHROMA = "chroma"
+    FAISS = "faiss"
+
+class EmbeddingType(str, Enum):
+    """Supported embedding model types."""
+    OPENAI = "openai"
+    HUGGINGFACE = "huggingface"
+    INSTRUCTOR = "instructor"
+    BEDROCK = "bedrock"
+
+class RagTool(Tool):
+    """Tool for performing RAG operations using LlamaIndex."""
+
+    name: str = "rag_tool"
+    description: str = (
+        "Retrieval Augmented Generation (RAG) tool for querying indexed documents "
+        "using vector stores and embedding models."
+    )
+    arguments: List[ToolArgument] = [
+        ToolArgument(
+            name="query",
+            arg_type="string",
+            description="Query string for searching the index",
+            required=True,
+            example="What is the main topic?",
+        ),
+    ]
+
+    def __init__(
+        self,
+        vector_store: str = "chroma",
+        embedding_model: str = "openai",
+        persist_dir: str = "./storage/rag",
+        document_paths: Optional[List[str]] = None,
+    ):
+        """Initialize the RAG tool with vector store, embedding model, and optional documents.
+
+        Args:
+            vector_store: Vector store type (chroma, faiss)
+            embedding_model: Embedding model type (openai, huggingface, instructor, bedrock)
+            persist_dir: Directory for persistence
+            document_paths: Optional list of paths to documents or directories to index
+        """
+        super().__init__()
+        self.persist_dir = os.path.abspath(persist_dir)
+        self.embed_model = self._setup_embedding_model(embedding_model)
+        self.vector_store = self._setup_vector_store(vector_store, self.persist_dir)
+        
+        # Configure llama-index settings with the embedding model
+        Settings.embed_model = self.embed_model
+        self.storage_context = StorageContext.from_defaults(vector_store=self.vector_store)
+        
+        # Initialize index
+        self.index = None
+        
+        # Check if we have documents to initialize with
+        if document_paths:
+            self._initialize_with_documents(document_paths)
+        else:
+            # Only try to load existing index if no documents were provided
+            index_exists = os.path.exists(os.path.join(self.persist_dir, "docstore.json"))
+            if index_exists:
+                try:
+                    self.index = load_index_from_storage(
+                        storage_context=self.storage_context,
+                    )
+                    logger.info(f"Loaded existing index from {self.persist_dir}")
+                except Exception as e:
+                    logger.error(f"Failed to load existing index: {str(e)}")
+                    self.index = None
+            else:
+                logger.warning("No existing index found and no documents provided")
+
+    def _initialize_with_documents(self, document_paths: List[str]) -> None:
+        """Initialize the index with the given documents.
+
+        Args:
+            document_paths: List of paths to documents or directories
+        """
+        try:
+            all_documents = []
+            for path in document_paths:
+                if not os.path.exists(path):
+                    logger.warning(f"Document path does not exist: {path}")
+                    continue
+                    
+                documents = SimpleDirectoryReader(
+                    input_files=[path] if os.path.isfile(path) else None,
+                    input_dir=path if os.path.isdir(path) else None,
+                ).load_data()
+                all_documents.extend(documents)
+            
+            if all_documents:
+                self.index = VectorStoreIndex.from_documents(
+                    all_documents,
+                    storage_context=self.storage_context,
+                )
+                # Persist the index after creation
+                self.storage_context.persist(persist_dir=self.persist_dir)
+                logger.info(f"Created and persisted new index with {len(all_documents)} documents")
+            else:
+                logger.warning("No valid documents found in provided paths")
+                
+        except Exception as e:
+            logger.error(f"Error initializing with documents: {str(e)}")
+            raise RuntimeError(f"Failed to initialize with documents: {str(e)}")
+
+    def _setup_embedding_model(self, model_type: str) -> Any:
+        """Set up the embedding model based on type.
+
+        Args:
+            model_type: Type of embedding model to use
+
+        Returns:
+            Configured embedding model instance
+        """
+        model_type = EmbeddingType(model_type.lower())
+        if model_type == EmbeddingType.OPENAI:
+            return OpenAIEmbedding()
+        elif model_type == EmbeddingType.HUGGINGFACE:
+            return HuggingFaceEmbedding()
+        elif model_type == EmbeddingType.INSTRUCTOR:
+            return InstructorEmbedding()
+        elif model_type == EmbeddingType.BEDROCK:
+            return BedrockEmbedding()
+        else:
+            raise ValueError(f"Unsupported embedding model type: {model_type}")
+
+    def _setup_vector_store(self, store_type: str, persist_dir: str) -> Any:
+        """Set up the vector store based on type.
+
+        Args:
+            store_type: Type of vector store to use
+            persist_dir: Directory for persistence
+
+        Returns:
+            Configured vector store instance
+        """
+        store_type = VectorStoreType(store_type.lower())
+        
+        # Ensure the persist directory exists
+        os.makedirs(persist_dir, exist_ok=True)
+        
+        if store_type == VectorStoreType.CHROMA:
+            # Use PersistentClient with explicit settings
+            chroma_persist_dir = os.path.join(persist_dir, "chroma")
+            os.makedirs(chroma_persist_dir, exist_ok=True)
+            
+            chroma_client = chromadb.PersistentClient(
+                path=chroma_persist_dir,
+            )
+            collection = chroma_client.create_collection(
+                name="default_collection",
+                get_or_create=True
+            )
+            return ChromaVectorStore(
+                chroma_collection=collection,
+            )
+        elif store_type == VectorStoreType.FAISS:
+            return FaissVectorStore()
+        else:
+            raise ValueError(f"Unsupported vector store type: {store_type}")
+
+    def add_documents(self, document_path: str) -> bool:
+        """Add documents to the RAG system.
+
+        Args:
+            document_path: Path to document or directory of documents
+
+        Returns:
+            bool: True if documents were added successfully
+        """
+        try:
+            if not os.path.exists(document_path):
+                logger.error(f"Document path does not exist: {document_path}")
+                return False
+
+            documents = SimpleDirectoryReader(
+                input_files=[document_path] if os.path.isfile(document_path) else None,
+                input_dir=document_path if os.path.isdir(document_path) else None,
+            ).load_data()
+
+            # Create index with configured settings and storage context
+            self.index = VectorStoreIndex.from_documents(
+                documents,
+                storage_context=self.storage_context,
+            )
+            
+            return True
+        except Exception as e:
+            logger.error(f"Error adding documents: {str(e)}")
+            return False
+
+    def execute(self, query: str) -> str:
+        """Execute a query against the indexed documents.
+
+        Args:
+            query: Query string for searching
+
+        Returns:
+            Query response
+
+        Raises:
+            ValueError: If no index is available
+        """
+        try:
+            if not self.index:
+                raise ValueError("No index available. Please add documents first using add_documents()")
+
+            # Query the index
+            query_engine = self.index.as_query_engine()
+            response = query_engine.query(query)
+            return str(response)
+
+        except Exception as e:
+            logger.error(f"Error in RAG query: {str(e)}")
+            raise RuntimeError(f"Query failed: {str(e)}")
+
+
+if __name__ == "__main__":
+    # Example usage
+    tool = RagTool(
+        vector_store="chroma",
+        embedding_model="openai",
+        persist_dir="./storage/rag",
+        document_paths=[
+            "./docs/file1.pdf",
+            "./docs/directory1"
+        ]
+    )
+    
+    # Query
+    print(tool.execute("What is the main topic?"))

quantalogic/tools/read_html_tool.py

@@ -32,56 +32,48 @@ USER_AGENTS = [
     "Mozilla/5.0 (Macintosh; Intel Mac OS X 10.15; rv:89.0) Gecko/20100101 Firefox/89.0",
     # Safari on macOS
     "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/605.1.15 (KHTML, like Gecko)"
-    " Version/14.1.1 Safari/605.1.15"
+    " Version/14.1.1 Safari/605.1.15",
 ]
 
 # Additional headers to mimic real browser requests
 ADDITIONAL_HEADERS = {
-    "Accept": "text/html,application/xhtml+xml,application/xml;"
-              "q=0.9,image/webp,*/*;q=0.8",
+    "Accept": "text/html,application/xhtml+xml,application/xml;" "q=0.9,image/webp,*/*;q=0.8",
     "Accept-Language": "en-US,en;q=0.5",
     "Upgrade-Insecure-Requests": "1",
     "DNT": "1",  # Do Not Track
     "Connection": "keep-alive",
-    "Cache-Control": "max-age=0"
+    "Cache-Control": "max-age=0",
 }
 
+
 class ReadHTMLTool(Tool):
     """Tool for reading HTML content from files or URLs in specified line ranges."""
 
     class Arguments(BaseModel):
         source: str = Field(
-            ...,
-            description="The file path or URL to read HTML from",
-            example="https://example.com or ./example.html"
+            ..., description="The file path or URL to read HTML from", example="https://example.com or ./example.html"
         )
         convert: Optional[str] = Field(
             "text",
             description="Convert input to 'text' (Markdown) or 'html' no conversion. Default is 'text'",
-            example="'text' or 'html'"
+            example="'text' or 'html'",
         )
         line_start: Optional[int] = Field(
-            1,
-            description="The starting line number (1-based index). Default: 1",
-            ge=1,
-            example="1"
+            1, description="The starting line number (1-based index). Default: 1", ge=1, example="1"
         )
         line_end: Optional[int] = Field(
-            300,
-            description="The ending line number (1-based index). Default: 300",
-            ge=1,
-            example="300"
+            300, description="The ending line number (1-based index). Default: 300", ge=1, example="300"
         )
 
-        @field_validator('convert')
+        @field_validator("convert")
         def validate_convert(cls, v):
             if v not in ["text", "html"]:
                 raise ValueError("Convert must be either 'text' or 'html'")
             return v
 
-        @field_validator('line_end')
+        @field_validator("line_end")
         def validate_line_end(cls, v, values):
-            if 'line_start' in values and v < values['line_start']:
+            if "line_start" in values and v < values["line_start"]:
                 raise ValueError("line_end must be greater than or equal to line_start")
             return v
 
@@ -97,13 +89,13 @@ class ReadHTMLTool(Tool):
             arg_type="string",
             description="The file path or URL to read HTML from",
             required=True,
-            example="https://example.com or ./example.html"
+            example="https://example.com or ./example.html",
         ),
         ToolArgument(
             name="convert",
             arg_type="string",
             description="Convert input to 'text' (Markdown) or 'html'. Default is 'text'",
-            default='text',
+            default="text",
             required=False,
             example="'text' or 'html'",
         ),
@@ -113,7 +105,7 @@ class ReadHTMLTool(Tool):
             description="The starting line number (1-based index). Default: 1",
             required=False,
             example="1",
-            default="1"
+            default="1",
         ),
         ToolArgument(
             name="line_end",
@@ -121,8 +113,8 @@ class ReadHTMLTool(Tool):
             description="The ending line number (1-based index). Default: 300",
             required=False,
             example="300",
-            default="300"
-        )
+            default="300",
+        ),
     ]
 
     def validate_source(self, source: str) -> bool:
@@ -140,7 +132,7 @@ class ReadHTMLTool(Tool):
     def read_from_file(self, file_path: str) -> str:
         """Read HTML content from a file."""
         try:
-            with open(file_path, encoding='utf-8') as file:
+            with open(file_path, encoding="utf-8") as file:
                 return file.read()
         except Exception as e:
             logger.error(f"Error reading file: {e}")
@@ -157,12 +149,7 @@ class ReadHTMLTool(Tool):
             time.sleep(random.uniform(0.5, 2.0))
 
             # Use a timeout to prevent hanging
-            response = requests.get(
-                url,
-                headers=headers,
-                timeout=10,
-                allow_redirects=True
-            )
+            response = requests.get(url, headers=headers, timeout=10, allow_redirects=True)
             response.raise_for_status()
             return response.text
         except requests.RequestException as e:
@@ -172,7 +159,7 @@ class ReadHTMLTool(Tool):
     def parse_html(self, html_content: str) -> BeautifulSoup:
         """Parse HTML content using BeautifulSoup."""
         try:
-            return BeautifulSoup(html_content, 'html.parser')
+            return BeautifulSoup(html_content, "html.parser")
         except Exception as e:
             logger.error(f"Error parsing HTML: {e}")
             raise ValueError(f"Error parsing HTML: {e}")
@@ -207,7 +194,7 @@ class ReadHTMLTool(Tool):
         if convert_type == "html":
             # Ensure content is valid HTML
             try:
-                soup = BeautifulSoup(content, 'html.parser')
+                soup = BeautifulSoup(content, "html.parser")
                 return soup.prettify()
             except Exception as e:
                 logger.error(f"Error prettifying HTML: {e}")
@@ -215,8 +202,7 @@ class ReadHTMLTool(Tool):
 
         return content
 
-    def execute(self, source: str, convert: Optional[str] = 'text',
-                line_start: int = 1, line_end: int = 300) -> str:
+    def execute(self, source: str, convert: Optional[str] = "text", line_start: int = 1, line_end: int = 300) -> str:
         """Execute the tool to read and parse HTML content in specified line ranges."""
         logger.debug(f"Executing read_html_tool with source: {source}")
 
@@ -242,7 +228,7 @@ class ReadHTMLTool(Tool):
             adjusted_end_line = min(line_end, total_lines)
 
             # Step 5: Slice lines based on line_start and adjusted_end_line
-            sliced_lines = lines[line_start - 1: adjusted_end_line]
+            sliced_lines = lines[line_start - 1 : adjusted_end_line]
             sliced_content = "\n".join(sliced_lines)
 
             # Step 6: Calculate actual_end_line based on lines returned
@@ -268,7 +254,7 @@ class ReadHTMLTool(Tool):
                 f"Is Last Block: {'Yes' if is_last_block else 'No'}",
                 "==== Content ====",
                 sliced_content,
-                "==== End of Block ===="
+                "==== End of Block ====",
             ]
 
             return "\n".join(result)
@@ -302,4 +288,4 @@ if __name__ == "__main__":
         else:
             print("No local test file found.")
     except Exception as e:
-        print(f"Local File Test Failed: {e}")
+        print(f"Local File Test Failed: {e}")

quantalogic/tools/replace_in_file_tool.py

@@ -64,7 +64,7 @@ class ReplaceInFileTool(Tool):
         "Returns the updated content or an error."
         "⚠️ THIS TOOL MUST BE USED IN PRIORITY TO UPDATE AN EXISTING FILE."
     )
-    need_validation: bool = True
+    need_validation: bool = False
 
     SIMILARITY_THRESHOLD: float = 0.85
 
@@ -126,12 +126,12 @@ class ReplaceInFileTool(Tool):
 
     def normalize_whitespace(self, text: str) -> str:
         """Normalize leading whitespace by converting tabs to spaces."""
-        return '\n'.join([self._normalize_line(line) for line in text.split('\n')])
+        return "\n".join([self._normalize_line(line) for line in text.split("\n")])
 
     def _normalize_line(self, line: str) -> str:
         """Normalize leading whitespace in a single line."""
         leading_ws = len(line) - len(line.lstrip())
-        return line.replace('\t', '    ', leading_ws)  # Convert tabs to 4 spaces only in leading whitespace
+        return line.replace("\t", "    ", leading_ws)  # Convert tabs to 4 spaces only in leading whitespace
 
     def parse_diff(self, diff: str) -> list[SearchReplaceBlock]:
         """Parses the diff string into a list of SearchReplaceBlock instances."""
@@ -279,25 +279,25 @@ class ReplaceInFileTool(Tool):
     def find_similar_match(self, search: str, content: str) -> Tuple[float, str]:
         """Finds the most similar substring in content compared to search with whitespace normalization."""
         norm_search = self.normalize_whitespace(search)
-        content_lines = content.split('\n')
+        content_lines = content.split("\n")
         norm_content = self.normalize_whitespace(content)
-        norm_content_lines = norm_content.split('\n')
+        norm_content_lines = norm_content.split("\n")
 
-        if len(norm_content_lines) < len(norm_search.split('\n')):
+        if len(norm_content_lines) < len(norm_search.split("\n")):
             return 0.0, ""
 
         max_similarity = 0.0
         best_match = ""
-        search_line_count = len(norm_search.split('\n'))
+        search_line_count = len(norm_search.split("\n"))
 
         for i in range(len(norm_content_lines) - search_line_count + 1):
-            candidate_norm = '\n'.join(norm_content_lines[i:i+search_line_count])
+            candidate_norm = "\n".join(norm_content_lines[i : i + search_line_count])
             similarity = difflib.SequenceMatcher(None, norm_search, candidate_norm).ratio()
 
             if similarity > max_similarity:
                 max_similarity = similarity
                 # Get original lines (non-normalized) for accurate replacement
-                best_match = '\n'.join(content_lines[i:i+search_line_count])
+                best_match = "\n".join(content_lines[i : i + search_line_count])
 
         return max_similarity, best_match
 
@@ -308,4 +308,4 @@ class ReplaceInFileTool(Tool):
 
 if __name__ == "__main__":
     tool = ReplaceInFileTool()
-    print(tool.to_markdown())
+    print(tool.to_markdown())

quantalogic/tools/safe_python_interpreter_tool.py

@@ -29,17 +29,12 @@ class SafePythonInterpreterTool(Tool):
     A tool to safely execute Python code while only allowing a specific set
     of modules as defined in `allowed_modules`.
     """
+
     # Allowed modules must be provided during initialization.
-    allowed_modules: List[str] = Field(
-        ...,
-        description="List of Python module names allowed for code execution."
-    )
+    allowed_modules: List[str] = Field(..., description="List of Python module names allowed for code execution.")
     # Additional fields to support the Tool API.
     code: str | None = None  # Provided at runtime via kwargs.
-    time_limit: int = Field(
-        default=60,
-        description="Maximum execution time (in seconds) for running the Python code."
-    )
+    time_limit: int = Field(default=60, description="Maximum execution time (in seconds) for running the Python code.")
     # Define tool arguments so that they appear in the tool's markdown description.
     arguments: list[ToolArgument] = [
         ToolArgument(
@@ -59,7 +54,7 @@ def transform_array(x):
 array_input = np.array([1, 4, 9, 16, 25])
 result = transform_array(array_input)
 result
-            """.strip()
+            """.strip(),
         ),
         ToolArgument(
             name="time_limit",
@@ -67,8 +62,8 @@ result
             description="The execution timeout (in seconds).",
             required=False,
             default="60",
-            example="60"
-        )
+            example="60",
+        ),
     ]
     name: Literal["safe_python_interpreter"] = "safe_python_interpreter"
     description: str | None = None
@@ -114,18 +109,12 @@ result
         def run_interpreter() -> Any:
             logger.debug("Starting interpretation of code.")
             import ast  # new import for AST processing
+
             # Delegate to monkeypatched interpret_code if available.
             if interpret_code.__module__ != "quantalogic.utils.python_interpreter":
                 return interpret_code(code, self.allowed_modules)
             # Build safe globals with only allowed modules and minimal builtins.
-            safe_globals = {
-                "__builtins__": {
-                    "range": range,
-                    "len": len,
-                    "print": print,
-                    "__import__": __import__
-                }
-            }
+            safe_globals = {"__builtins__": {"range": range, "len": len, "print": print, "__import__": __import__}}
             for mod in self.allowed_modules:
                 safe_globals[mod] = __import__(mod)
             local_vars = {}
@@ -139,10 +128,7 @@ result
                 tree = ast.parse(code)
                 if tree.body and isinstance(tree.body[-1], ast.Expr):
                     last_expr = tree.body.pop()
-                    assign = ast.Assign(
-                        targets=[ast.Name(id="_result", ctx=ast.Store())],
-                        value=last_expr.value
-                    )
+                    assign = ast.Assign(targets=[ast.Name(id="_result", ctx=ast.Store())], value=last_expr.value)
                     assign = ast.copy_location(assign, last_expr)
                     tree.body.append(assign)
                     fixed_tree = ast.fix_missing_locations(tree)
@@ -210,4 +196,4 @@ result
         print("Interpreter Output:")
         print(output)
     except Exception as e:
-        print(f"An error occurred during interpretation: {e}")
+        print(f"An error occurred during interpretation: {e}")

quantalogic/tools/search_definition_names.py

@@ -69,7 +69,7 @@ class SearchDefinitionNames(Tool):
             description="The page number to retrieve (1-based index).",
             required=False,
             example="1",
-            default="1"
+            default="1",
         ),
         ToolArgument(
             name="page_size",
@@ -77,7 +77,7 @@ class SearchDefinitionNames(Tool):
             description="The number of results per page (default: 10).",
             required=False,
             example="10",
-            default="10"
+            default="10",
         ),
     ]
 

quantalogic/tools/sequence_tool.py

@@ -61,9 +61,9 @@ class SequenceTool(Tool):
       wraps the tool calls in a single <sequence> element or is the raw collection of tool calls.
       The tag for each tool call must be one of the authorized tool names.
       Authorized tools: {authorized_list}
-      
+
       For each tool call, include the required parameters as sub-elements. The syntax is as follows:
-      
+
       <sequence>
           <tool_name>
               <param1>value1</param1>
@@ -72,7 +72,7 @@ class SequenceTool(Tool):
           </tool_name>
           ...
       </sequence>
-      
+
     Objective:
       Execute each tool call in order—even if some calls fail—and return an XML report that includes
       the order of execution, provided parameters, execution status (success/failure), and the output or
@@ -149,11 +149,7 @@ class SequenceTool(Tool):
 
         # Try to locate <sequence> elements in the input.
         all_elements = xml_parser._find_all_elements(sequence_xml)
-        sequence_elems = [
-            (tag, content)
-            for tag, content in all_elements
-            if tag.strip().lower() == "sequence"
-        ]
+        sequence_elems = [(tag, content) for tag, content in all_elements if tag.strip().lower() == "sequence"]
         if len(sequence_elems) > 1:
             raise ValueError("Input XML must contain exactly one <sequence> element.")
         elif len(sequence_elems) == 1:
@@ -225,20 +221,15 @@ if __name__ == "__main__":
         def __init__(self, **data: Any):
             data.setdefault("name", "write_file_tool")
             data.setdefault("description", "Writes content to a file at the specified path.")
-            data.setdefault("arguments", [
-                ToolArgument(
-                    name="file_path",
-                    arg_type="string",
-                    required=True,
-                    description="Path to the file"
-                ),
-                ToolArgument(
-                    name="content",
-                    arg_type="string",
-                    required=True,
-                    description="Content to write to the file"
-                )
-            ])
+            data.setdefault(
+                "arguments",
+                [
+                    ToolArgument(name="file_path", arg_type="string", required=True, description="Path to the file"),
+                    ToolArgument(
+                        name="content", arg_type="string", required=True, description="Content to write to the file"
+                    ),
+                ],
+            )
             super().__init__(**data)
 
         def execute(self, **kwargs) -> str:
@@ -282,4 +273,4 @@ And hope finds its own sweet way.
         result_xml = sequence_tool.execute(sequence=xml_sequence)
         print("SequenceTool result:\n", result_xml)
     except Exception as e:
-        print("SequenceTool error:\n", str(e))
+        print("SequenceTool error:\n", str(e))