MindsDB 25.9.2.0a1__py3-none-any.whl → 25.10.0rc1__py3-none-any.whl

mindsdb/integrations/utilities/files/file_reader.py

@@ -1,17 +1,16 @@
-from dataclasses import dataclass, astuple
-import traceback
-import json
 import csv
-from io import BytesIO, StringIO, IOBase
-from pathlib import Path
+import json
 import codecs
+from io import BytesIO, StringIO, IOBase
 from typing import List, Generator
+from pathlib import Path
+from dataclasses import dataclass, astuple
 
 import filetype
 import pandas as pd
 from charset_normalizer import from_bytes
-from mindsdb.interfaces.knowledge_base.preprocessing.text_splitter import TextSplitter
 
+from mindsdb.interfaces.knowledge_base.preprocessing.text_splitter import TextSplitter
 from mindsdb.utilities import log
 
 logger = log.getLogger(__name__)
@@ -76,7 +75,7 @@ def decode(file_obj: IOBase) -> StringIO:
 
                 data_str = StringIO(byte_str.decode(encoding, errors))
     except Exception as e:
-        logger.error(traceback.format_exc())
+        logger.exception("Error during file decode:")
         raise FileProcessingError("Could not load into string") from e
 
     return data_str

mindsdb/integrations/utilities/handlers/auth_utilities/snowflake/__init__.py

@@ -0,0 +1 @@
+from .snowflake_jwt_gen import get_validated_jwt as get_validated_jwt

mindsdb/integrations/utilities/handlers/auth_utilities/snowflake/snowflake_jwt_gen.py

@@ -0,0 +1,151 @@
+# Based on https://docs.snowflake.com/en/developer-guide/sql-api/authenticating
+
+import time
+import base64
+import hashlib
+import logging
+from datetime import timedelta, timezone, datetime
+
+from cryptography.hazmat.primitives.serialization import load_pem_private_key
+from cryptography.hazmat.primitives.serialization import Encoding
+from cryptography.hazmat.primitives.serialization import PublicFormat
+from cryptography.hazmat.backends import default_backend
+import jwt
+
+logger = logging.getLogger(__name__)
+
+ISSUER = "iss"
+EXPIRE_TIME = "exp"
+ISSUE_TIME = "iat"
+SUBJECT = "sub"
+
+
+class JWTGenerator(object):
+    """
+    Creates and signs a JWT with the specified private key file, username, and account identifier. The JWTGenerator keeps the
+    generated token and only regenerates the token if a specified period of time has passed.
+    """
+
+    LIFETIME = timedelta(minutes=60)  # The tokens will have a 59 minute lifetime
+    ALGORITHM = "RS256"  # Tokens will be generated using RSA with SHA256
+
+    def __init__(self, account: str, user: str, private_key: str, lifetime: timedelta = LIFETIME):
+        """
+        __init__ creates an object that generates JWTs for the specified user, account identifier, and private key.
+        :param account: Your Snowflake account identifier. See https://docs.snowflake.com/en/user-guide/admin-account-identifier.html. Note that if you are using the account locator, exclude any region information from the account locator.
+        :param user: The Snowflake username.
+        :param private_key: The private key file used for signing the JWTs.
+        :param lifetime: The number of minutes (as a timedelta) during which the key will be valid.
+        """
+
+        logger.info(
+            """Creating JWTGenerator with arguments
+            account : %s, user : %s, lifetime : %s""",
+            account,
+            user,
+            lifetime,
+        )
+
+        # Construct the fully qualified name of the user in uppercase.
+        self.account = self.prepare_account_name_for_jwt(account)
+        self.user = user.upper()
+        self.qualified_username = self.account + "." + self.user
+
+        self.lifetime = lifetime
+        self.renew_time = datetime.now(timezone.utc)
+        self.token = None
+
+        self.private_key = load_pem_private_key(private_key.encode(), None, default_backend())
+
+    def prepare_account_name_for_jwt(self, raw_account: str) -> str:
+        """
+        Prepare the account identifier for use in the JWT.
+        For the JWT, the account identifier must not include the subdomain or any region or cloud provider information.
+        :param raw_account: The specified account identifier.
+        :return: The account identifier in a form that can be used to generate JWT.
+        """
+        account = raw_account
+        if ".global" not in account:
+            # Handle the general case.
+            idx = account.find(".")
+            if idx > 0:
+                account = account[0:idx]
+        else:
+            # Handle the replication case.
+            idx = account.find("-")
+            if idx > 0:
+                account = account[0:idx]
+        # Use uppercase for the account identifier.
+        return account.upper()
+
+    def get_token(self) -> str:
+        """
+        Generates a new JWT.
+        :return: the new token
+        """
+        now = datetime.now(timezone.utc)  # Fetch the current time
+
+        # Prepare the fields for the payload.
+        # Generate the public key fingerprint for the issuer in the payload.
+        public_key_fp = self.calculate_public_key_fingerprint(self.private_key)
+
+        # Create our payload
+        payload = {
+            # Set the issuer to the fully qualified username concatenated with the public key fingerprint.
+            ISSUER: self.qualified_username + "." + public_key_fp,
+            # Set the subject to the fully qualified username.
+            SUBJECT: self.qualified_username,
+            # Set the issue time to now.
+            ISSUE_TIME: now,
+            # Set the expiration time, based on the lifetime specified for this object.
+            EXPIRE_TIME: now + self.lifetime,
+        }
+
+        # Regenerate the actual token
+        token = jwt.encode(payload, key=self.private_key, algorithm=JWTGenerator.ALGORITHM)
+        # If you are using a version of PyJWT prior to 2.0, jwt.encode returns a byte string, rather than a string.
+        # If the token is a byte string, convert it to a string.
+        if isinstance(token, bytes):
+            token = token.decode("utf-8")
+        self.token = token
+
+        return self.token
+
+    def calculate_public_key_fingerprint(self, private_key: str) -> str:
+        """
+        Given a private key in PEM format, return the public key fingerprint.
+        :param private_key: private key string
+        :return: public key fingerprint
+        """
+        # Get the raw bytes of public key.
+        public_key_raw = private_key.public_key().public_bytes(Encoding.DER, PublicFormat.SubjectPublicKeyInfo)
+
+        # Get the sha256 hash of the raw bytes.
+        sha256hash = hashlib.sha256()
+        sha256hash.update(public_key_raw)
+
+        # Base64-encode the value and prepend the prefix 'SHA256:'.
+        public_key_fp = "SHA256:" + base64.b64encode(sha256hash.digest()).decode("utf-8")
+        logger.info("Public key fingerprint is %s", public_key_fp)
+
+        return public_key_fp
+
+
+def get_validated_jwt(token: str, account: str, user: str, private_key: str) -> str:
+    try:
+        content = jwt.decode(token, algorithms=[JWTGenerator.ALGORITHM], options={"verify_signature": False})
+
+        expired = content.get("exp", 0)
+        # add 5 seconds before limit
+        if expired - 5 > time.time():
+            # keep the same
+            return token
+
+    except jwt.DecodeError:
+        # wrong key
+        ...
+
+    # generate new token
+    if private_key is None:
+        raise ValueError("Private key is missing")
+    return JWTGenerator(account, user, private_key).get_token()

mindsdb/integrations/utilities/rag/config_loader.py

@@ -1,17 +1,26 @@
 """Utility functions for RAG pipeline configuration"""
+
 from typing import Dict, Any, Optional
 
 from mindsdb.utilities.log import getLogger
 from mindsdb.integrations.utilities.rag.settings import (
-    RetrieverType, MultiVectorRetrieverMode, SearchType,
-    SearchKwargs, SummarizationConfig, VectorStoreConfig,
-    RerankerConfig, RAGPipelineModel, DEFAULT_COLLECTION_NAME
+    RetrieverType,
+    MultiVectorRetrieverMode,
+    SearchType,
+    SearchKwargs,
+    SummarizationConfig,
+    VectorStoreConfig,
+    RerankerConfig,
+    RAGPipelineModel,
+    DEFAULT_COLLECTION_NAME,
 )
 
 logger = getLogger(__name__)
 
 
-def load_rag_config(base_config: Dict[str, Any], kb_params: Optional[Dict[str, Any]] = None, embedding_model: Any = None) -> RAGPipelineModel:
+def load_rag_config(
+    base_config: Dict[str, Any], kb_params: Optional[Dict[str, Any]] = None, embedding_model: Any = None
+) -> RAGPipelineModel:
     """
     Load and validate RAG configuration parameters. This function handles the conversion of configuration
     parameters into their appropriate types and ensures all required settings are properly configured.
@@ -37,41 +46,43 @@ def load_rag_config(base_config: Dict[str, Any], kb_params: Optional[Dict[str, A
 
     # Set embedding model if provided
     if embedding_model is not None:
-        rag_params['embedding_model'] = embedding_model
+        rag_params["embedding_model"] = embedding_model
 
     # Handle enums and type conversions
-    if 'retriever_type' in rag_params:
-        rag_params['retriever_type'] = RetrieverType(rag_params['retriever_type'])
-    if 'multi_retriever_mode' in rag_params:
-        rag_params['multi_retriever_mode'] = MultiVectorRetrieverMode(rag_params['multi_retriever_mode'])
-    if 'search_type' in rag_params:
-        rag_params['search_type'] = SearchType(rag_params['search_type'])
+    if "retriever_type" in rag_params:
+        rag_params["retriever_type"] = RetrieverType(rag_params["retriever_type"])
+    if "multi_retriever_mode" in rag_params:
+        rag_params["multi_retriever_mode"] = MultiVectorRetrieverMode(rag_params["multi_retriever_mode"])
+    if "search_type" in rag_params:
+        rag_params["search_type"] = SearchType(rag_params["search_type"])
 
     # Handle search kwargs if present
-    if 'search_kwargs' in rag_params and isinstance(rag_params['search_kwargs'], dict):
-        rag_params['search_kwargs'] = SearchKwargs(**rag_params['search_kwargs'])
+    if "search_kwargs" in rag_params and isinstance(rag_params["search_kwargs"], dict):
+        rag_params["search_kwargs"] = SearchKwargs(**rag_params["search_kwargs"])
 
     # Handle summarization config if present
-    summarization_config = rag_params.get('summarization_config')
+    summarization_config = rag_params.get("summarization_config")
     if summarization_config is not None and isinstance(summarization_config, dict):
-        rag_params['summarization_config'] = SummarizationConfig(**summarization_config)
+        rag_params["summarization_config"] = SummarizationConfig(**summarization_config)
 
     # Handle vector store config
-    if 'vector_store_config' in rag_params:
-        if isinstance(rag_params['vector_store_config'], dict):
-            rag_params['vector_store_config'] = VectorStoreConfig(**rag_params['vector_store_config'])
+    if "vector_store_config" in rag_params:
+        if isinstance(rag_params["vector_store_config"], dict):
+            rag_params["vector_store_config"] = VectorStoreConfig(**rag_params["vector_store_config"])
     else:
-        rag_params['vector_store_config'] = {}
-        logger.warning(f'No collection_name specified for the retrieval tool, '
-                       f"using default collection_name: '{DEFAULT_COLLECTION_NAME}'"
-                       f'\nWarning: If this collection does not exist, no data will be retrieved')
+        rag_params["vector_store_config"] = {}
+        logger.warning(
+            f"No collection_name specified for the retrieval tool, "
+            f"using default collection_name: '{DEFAULT_COLLECTION_NAME}'"
+            f"\nWarning: If this collection does not exist, no data will be retrieved"
+        )
 
-    if 'reranker_config' in rag_params:
-        rag_params['reranker_config'] = RerankerConfig(**rag_params['reranker_config'])
+    if "reranker_config" in rag_params:
+        rag_params["reranker_config"] = RerankerConfig(**rag_params["reranker_config"])
 
     # Convert to RAGPipelineModel with validation
     try:
         return RAGPipelineModel(**rag_params)
     except Exception as e:
-        logger.error(f"Invalid RAG configuration: {str(e)}")
-        raise ValueError(f"Configuration validation failed: {str(e)}")
+        logger.exception("Invalid RAG configuration:")
+        raise ValueError(f"Configuration validation failed: {str(e)}") from e

mindsdb/integrations/utilities/rag/rerankers/base_reranker.py

@@ -7,18 +7,35 @@ import math
 import os
 import random
 from abc import ABC
-from textwrap import dedent
 from typing import Any, List, Optional, Tuple
 
 from openai import AsyncOpenAI, AsyncAzureOpenAI
 from pydantic import BaseModel
 
-from mindsdb.integrations.utilities.rag.settings import DEFAULT_RERANKING_MODEL, DEFAULT_LLM_ENDPOINT
+from mindsdb.integrations.utilities.rag.settings import (
+    DEFAULT_RERANKING_MODEL,
+    DEFAULT_LLM_ENDPOINT,
+    DEFAULT_RERANKER_N,
+    DEFAULT_RERANKER_LOGPROBS,
+    DEFAULT_RERANKER_TOP_LOGPROBS,
+    DEFAULT_RERANKER_MAX_TOKENS,
+    DEFAULT_VALID_CLASS_TOKENS,
+)
 from mindsdb.integrations.libs.base import BaseMLEngine
 
 log = logging.getLogger(__name__)
 
 
+def get_event_loop():
+    try:
+        loop = asyncio.get_running_loop()
+    except RuntimeError:
+        # If no running loop exists, create a new one
+        loop = asyncio.new_event_loop()
+        asyncio.set_event_loop(loop)
+    return loop
+
+
 class BaseLLMReranker(BaseModel, ABC):
     filtering_threshold: float = 0.0  # Default threshold for filtering
     provider: str = "openai"
@@ -38,6 +55,11 @@ class BaseLLMReranker(BaseModel, ABC):
     request_timeout: float = 20.0  # Timeout for API requests
     early_stop: bool = True  # Whether to enable early stopping
     early_stop_threshold: float = 0.8  # Confidence threshold for early stopping
+    n: int = DEFAULT_RERANKER_N  # Number of completions to generate
+    logprobs: bool = DEFAULT_RERANKER_LOGPROBS  # Whether to include log probabilities
+    top_logprobs: int = DEFAULT_RERANKER_TOP_LOGPROBS  # Number of top log probabilities to include
+    max_tokens: int = DEFAULT_RERANKER_MAX_TOKENS  # Maximum tokens to generate
+    valid_class_tokens: List[str] = DEFAULT_VALID_CLASS_TOKENS
 
     class Config:
         arbitrary_types_allowed = True
@@ -61,7 +83,12 @@ class BaseLLMReranker(BaseModel, ABC):
                     timeout=self.request_timeout,
                     max_retries=2,
                 )
-            elif self.provider == "openai":
+            elif self.provider in ("openai", "ollama"):
+                if self.provider == "ollama":
+                    self.method = "no-logprobs"
+                    if self.api_key is None:
+                        self.api_key = "n/a"
+
                 api_key_var: str = "OPENAI_API_KEY"
                 openai_api_key = self.api_key or os.getenv(api_key_var)
                 if not openai_api_key:
@@ -71,7 +98,6 @@ class BaseLLMReranker(BaseModel, ABC):
                 self.client = AsyncOpenAI(
                     api_key=openai_api_key, base_url=base_url, timeout=self.request_timeout, max_retries=2
                 )
-
             else:
                 # try to use litellm
                 from mindsdb.api.executor.controllers.session_controller import SessionController
@@ -86,7 +112,7 @@ class BaseLLMReranker(BaseModel, ABC):
                 self.method = "no-logprobs"
 
     async def _call_llm(self, messages):
-        if self.provider in ("azure_openai", "openai"):
+        if self.provider in ("azure_openai", "openai", "ollama"):
             return await self.client.chat.completions.create(
                 model=self.model,
                 messages=messages,
@@ -121,7 +147,7 @@ class BaseLLMReranker(BaseModel, ABC):
             for idx, result in enumerate(results):
                 if isinstance(result, Exception):
                     log.error(f"Error processing document {i + idx}: {str(result)}")
-                    raise RuntimeError(f"Error during reranking: {result}")
+                    raise RuntimeError(f"Error during reranking: {result}") from result
 
                 score = result["relevance_score"]
 
@@ -142,7 +168,7 @@ class BaseLLMReranker(BaseModel, ABC):
                         return ranked_results
                 except Exception as e:
                     # Don't let early stopping errors stop the whole process
-                    log.warning(f"Error in early stopping check: {str(e)}")
+                    log.warning(f"Error in early stopping check: {e}")
 
         return ranked_results
 
@@ -204,13 +230,11 @@ class BaseLLMReranker(BaseModel, ABC):
         return rerank_data
 
     async def search_relevancy_no_logprob(self, query: str, document: str) -> Any:
-        prompt = dedent(
-            f"""
-            Score the relevance between search query and user message on scale between 0 and 100 per cents.
-            Consider semantic meaning, key concepts, and contextual relevance.
-            Return ONLY a numerical score between 0 and 100 per cents. No other text. Stop after sending a number
-            Search query: {query}
-        """
+        prompt = (
+            f"Score the relevance between search query and user message on scale between 0 and 100 per cents. "
+            f"Consider semantic meaning, key concepts, and contextual relevance. "
+            f"Return ONLY a numerical score between 0 and 100 per cents. No other text. Stop after sending a number. "
+            f"Search query: {query}"
         )
 
         response = await self._call_llm(
@@ -234,6 +258,28 @@ class BaseLLMReranker(BaseModel, ABC):
         return rerank_data
 
     async def search_relevancy_score(self, query: str, document: str) -> Any:
+        """
+        This method is used to score the relevance of a document to a query.
+
+        Args:
+            query: The query to score the relevance of.
+            document: The document to score the relevance of.
+
+        Returns:
+            A dictionary with the document and the relevance score.
+        """
+
+        log.debug("Start search_relevancy_score")
+        log.debug(f"Reranker query: {query[:5]}")
+        log.debug(f"Reranker document: {document[:50]}")
+        log.debug(f"Reranker model: {self.model}")
+        log.debug(f"Reranker temperature: {self.temperature}")
+        log.debug(f"Reranker n: {self.n}")
+        log.debug(f"Reranker logprobs: {self.logprobs}")
+        log.debug(f"Reranker top_logprobs: {self.top_logprobs}")
+        log.debug(f"Reranker max_tokens: {self.max_tokens}")
+        log.debug(f"Reranker valid_class_tokens: {self.valid_class_tokens}")
+
         response = await self.client.chat.completions.create(
             model=self.model,
             messages=[
@@ -306,17 +352,30 @@ class BaseLLMReranker(BaseModel, ABC):
                 },
             ],
             temperature=self.temperature,
-            n=1,
-            logprobs=True,
-            top_logprobs=4,
-            max_tokens=3,
+            n=self.n,
+            logprobs=self.logprobs,
+            top_logprobs=self.top_logprobs,
+            max_tokens=self.max_tokens,
         )
 
         # Extract response and logprobs
         token_logprobs = response.choices[0].logprobs.content
-        # Reconstruct the prediction and extract the top logprobs from the final token (e.g., "1")
-        final_token_logprob = token_logprobs[-1]
-        top_logprobs = final_token_logprob.top_logprobs
+
+        # Find the token that contains the class number
+        # Instead of just taking the last token, search for the actual class number token
+        class_token_logprob = None
+        for token_logprob in reversed(token_logprobs):
+            if token_logprob.token in self.valid_class_tokens:
+                class_token_logprob = token_logprob
+                break
+
+        # If we couldn't find a class token, fall back to the last non-empty token
+        if class_token_logprob is None:
+            log.warning("No class token logprob found, using the last token as fallback")
+            class_token_logprob = token_logprobs[-1]
+
+        top_logprobs = class_token_logprob.top_logprobs
+
         # Create a map of 'class_1' -> probability, using token combinations
         class_probs = {}
         for top_token in top_logprobs:
@@ -337,21 +396,15 @@ class BaseLLMReranker(BaseModel, ABC):
                 score = 0.0
 
         rerank_data = {"document": document, "relevance_score": score}
+        log.debug(f"Reranker score: {score}")
+        log.debug("End search_relevancy_score")
         return rerank_data
 
     def get_scores(self, query: str, documents: list[str]):
         query_document_pairs = [(query, doc) for doc in documents]
         # Create event loop and run async code
-        import asyncio
-
-        try:
-            loop = asyncio.get_running_loop()
-        except RuntimeError:
-            # If no running loop exists, create a new one
-            loop = asyncio.new_event_loop()
-            asyncio.set_event_loop(loop)
 
-        documents_and_scores = loop.run_until_complete(self._rank(query_document_pairs))
+        documents_and_scores = get_event_loop().run_until_complete(self._rank(query_document_pairs))
 
         scores = [score for _, score in documents_and_scores]
         return scores

mindsdb/integrations/utilities/rag/rerankers/reranker_compressor.py

@@ -36,7 +36,7 @@ class LLMReranker(BaseDocumentCompressor, BaseLLMReranker):
             return []
 
         # Stream reranking update.
-        dispatch_custom_event('rerank_begin', {'num_documents': len(documents)})
+        dispatch_custom_event("rerank_begin", {"num_documents": len(documents)})
 
         try:
             # Prepare query-document pairs
@@ -73,10 +73,10 @@ class LLMReranker(BaseDocumentCompressor, BaseLLMReranker):
             return filtered_docs
 
         except Exception as e:
-            error_msg = f"Error during async document compression: {str(e)}"
-            log.error(error_msg)
+            error_msg = "Error during async document compression:"
+            log.exception(error_msg)
             if callbacks:
-                await callbacks.on_retriever_error(error_msg)
+                await callbacks.on_retriever_error(f"{error_msg} {e}")
             return documents  # Return original documents on error
 
     def compress_documents(