MindsDB 25.1.2.0__py3-none-any.whl → 25.1.3.0__py3-none-any.whl

mindsdb/integrations/utilities/rag/settings.py

@@ -150,11 +150,118 @@ Here is the user input:
 {input}
 '''
 
+DEFAULT_QUESTION_REFORMULATION_TEMPLATE = """Given the original question and the retrieved context,
+analyze what additional information is needed for a complete, accurate answer.
+
+Original Question: {question}
+
+Retrieved Context:
+{context}
+
+Analysis Instructions:
+1. Evaluate Context Coverage:
+   - Identify key entities and concepts from the question
+   - Check for temporal information (dates, periods, sequences)
+   - Verify causal relationships are explained
+   - Confirm presence of requested quantitative data
+   - Assess if geographic or spatial context is sufficient
+
+2. Quality Assessment:
+   If the retrieved context is:
+   - Irrelevant or tangential
+   - Too general or vague
+   - Potentially contradictory
+   - Missing key perspectives
+   - Lacking proper evidence
+   Generate questions to address these specific gaps.
+
+3. Follow-up Question Requirements:
+   - Questions must directly contribute to answering the original query
+   - Break complex relationships into simpler, sequential steps
+   - Maintain specificity rather than broad inquiries
+   - Avoid questions answerable from existing context
+   - Ensure questions build on each other logically
+   - Limit questions to 150 characters each
+   - Each question must be self-contained
+   - Questions must end with a question mark
+
+4. Response Format:
+   - Return a JSON array of strings
+   - Use square brackets and double quotes
+   - Questions must be unique (no duplicates)
+   - If context is sufficient, return empty array []
+   - Maximum 3 follow-up questions
+   - Minimum length per question: 30 characters
+   - No null values or empty strings
+
+Example:
+Original: "How did the development of antibiotics affect military casualties in WWII?"
+
+Invalid responses:
+{'questions': ['What are antibiotics?']}  // Wrong format
+['What is WWII?']  // Too basic
+['How did it impact things?']  // Too vague
+['', 'Question 2']  // Contains empty string
+['Same question?', 'Same question?']  // Duplicate
+
+Valid response:
+["What were military casualty rates from infections before widespread antibiotic use in 1942?",
+ "How did penicillin availability change throughout different stages of WWII?",
+ "What were the primary battlefield infections treated with antibiotics during WWII?"]
+
+or [] if context fully answers the original question.
+
+Your task: Based on the analysis of the original question and context,
+output ONLY a JSON array of follow-up questions needed to provide a complete answer.
+If no additional information is needed, output an empty array [].
+
+Follow-up Questions:"""
+
+DEFAULT_QUERY_RETRY_PROMPT_TEMPLATE = '''
+{query}
+
+The {dialect} query above failed with the error message: {error}.
+
+<< TABLES YOU HAVE ACCESS TO >>
+1. {embeddings_table} - Contains document chunks, vector embeddings, and metadata for documents.
+
+Columns:
+```json
+{{
+    "id": {{
+        "type": "string",
+        "description": "Unique ID for this document chunk"
+    }},
+    "content": {{
+        "type": "string",
+        "description": "A document chunk (subset of the original document)"
+    }},
+    "embeddings": {{
+        "type": "vector",
+        "description": "Vector embeddings for the document chunk."
+    }},
+    "metadata": {{
+        "type": "jsonb",
+        "description": "Metadata for the document chunk."
+    }}
+}}
+
+{schema}
+
+Rewrite the query so it works.
+
+Output the final SQL query only.
+
+SQL Query:
+'''
+
+DEFAULT_NUM_QUERY_RETRIES = 2
+
 
 class LLMConfig(BaseModel):
     model_name: str = Field(default=DEFAULT_LLM_MODEL, description='LLM model to use for generation')
     provider: str = Field(default=DEFAULT_LLM_MODEL_PROVIDER, description='LLM model provider to use for generation')
-    params: Dict[str, Any] = {}
+    params: Dict[str, Any] = Field(default_factory=dict)
 
 
 class MultiVectorRetrieverMode(Enum):
@@ -183,17 +290,21 @@ class VectorStoreConfig(BaseModel):
     collection_name: str = DEFAULT_COLLECTION_NAME
     connection_string: str = None
     kb_table: Any = None
+    is_sparse: bool = False
+    vector_size: Optional[int] = None
 
     class Config:
         arbitrary_types_allowed = True
         extra = "forbid"
 
 
-class RetrieverType(Enum):
-    VECTOR_STORE = 'vector_store'
-    AUTO = 'auto'
-    MULTI = 'multi'
-    SQL = 'sql'
+class RetrieverType(str, Enum):
+    """Retriever type for RAG pipeline"""
+    VECTOR_STORE = "vector_store"
+    AUTO = "auto"
+    MULTI = "multi"
+    SQL = "sql"
+    MULTI_HOP = "multi_hop"
 
 
 class SearchType(Enum):
@@ -293,6 +404,14 @@ class SQLRetrieverConfig(BaseModel):
         default=DEFAULT_QUERY_CHECKER_PROMPT_TEMPLATE,
         description="Prompt template to use for double checking SQL queries before execution. Has 'query' and 'dialect' input variables."
     )
+    query_retry_template: str = Field(
+        default=DEFAULT_QUERY_RETRY_PROMPT_TEMPLATE,
+        description="Prompt template to rewrite SQL query that failed. Has 'dialect', 'query', and 'error' input variables."
+    )
+    num_retries: int = Field(
+        default=DEFAULT_NUM_QUERY_RETRIES,
+        description="How many times for an LLM to try rewriting a failed SQL query before using the fallback retriever."
+    )
     rewrite_prompt_template: str = Field(
         default=DEFAULT_SEMANTIC_PROMPT_TEMPLATE,
         description="Prompt template to rewrite user input to be better suited for retrieval. Has 'input' input variable."
@@ -336,6 +455,27 @@ class RerankerConfig(BaseModel):
     num_docs_to_keep: Optional[int] = None
 
 
+class MultiHopRetrieverConfig(BaseModel):
+    """Configuration for multi-hop retrieval"""
+    base_retriever_type: RetrieverType = Field(
+        default=RetrieverType.VECTOR_STORE,
+        description="Type of base retriever to use for multi-hop retrieval"
+    )
+    max_hops: int = Field(
+        default=3,
+        description="Maximum number of follow-up questions to generate",
+        ge=1
+    )
+    reformulation_template: str = Field(
+        default=DEFAULT_QUESTION_REFORMULATION_TEMPLATE,
+        description="Template for reformulating questions"
+    )
+    llm_config: LLMConfig = Field(
+        default_factory=LLMConfig,
+        description="LLM configuration to use for generating follow-up questions"
+    )
+
+
 class RAGPipelineModel(BaseModel):
     documents: Optional[List[Document]] = Field(
         default=None,
@@ -462,6 +602,20 @@ class RAGPipelineModel(BaseModel):
         description="Reranker configuration"
     )
 
+    multi_hop_config: Optional[MultiHopRetrieverConfig] = Field(
+        default=None,
+        description="Configuration for multi-hop retrieval. Required when retriever_type is MULTI_HOP."
+    )
+
+    @field_validator("multi_hop_config")
+    @classmethod
+    def validate_multi_hop_config(cls, v: Optional[MultiHopRetrieverConfig], info):
+        """Validate that multi_hop_config is set when using multi-hop retrieval."""
+        values = info.data
+        if values.get("retriever_type") == RetrieverType.MULTI_HOP and v is None:
+            raise ValueError("multi_hop_config must be set when using multi-hop retrieval")
+        return v
+
     class Config:
         arbitrary_types_allowed = True
         extra = "forbid"

mindsdb/integrations/utilities/sql_utils.py

@@ -178,7 +178,7 @@ def project_dataframe(df, targets, table_columns):
 
     # adapt column names to projection
     if len(df_col_rename) > 0:
-        df = df.rename(columns=df_col_rename)
+        df.rename(columns=df_col_rename, inplace=True)
     return df
 
 

mindsdb/interfaces/knowledge_base/controller.py

@@ -642,11 +642,13 @@ class KnowledgeBaseController:
             storage: Identifier,
             params: dict,
             preprocessing_config: Optional[dict] = None,
-            if_not_exists: bool = False,
+            if_not_exists: bool = False
     ) -> db.KnowledgeBase:
         """
         Add a new knowledge base to the database
         :param preprocessing_config: Optional preprocessing configuration to validate and store
+        :param is_sparse: Whether to use sparse vectors for embeddings
+        :param vector_size: Optional size specification for vectors, required when is_sparse=True
         """
         # Validate preprocessing config first if provided
         if preprocessing_config is not None:
@@ -654,6 +656,12 @@ class KnowledgeBaseController:
             params = params or {}
             params['preprocessing'] = preprocessing_config
 
+        # Check if vector_size is provided when using sparse vectors
+        is_sparse = params.get('is_sparse')
+        vector_size = params.get('vector_size')
+        if is_sparse and vector_size is None:
+            raise ValueError("vector_size is required when is_sparse=True")
+
         # get project id
         project = self.session.database_controller.get_project(project_name)
         project_id = project.id
@@ -693,7 +701,20 @@ class KnowledgeBaseController:
             cloud_pg_vector = os.environ.get('KB_PGVECTOR_URL')
             if cloud_pg_vector:
                 vector_table_name = name
-                vector_db_name = self._create_persistent_pgvector()
+                # Add sparse vector support for pgvector
+                vector_db_params = {}
+                # Check both explicit parameter and model configuration
+                is_sparse = is_sparse or model_record.learn_args.get('using', {}).get('sparse')
+                if is_sparse:
+                    vector_db_params['is_sparse'] = True
+                    if vector_size is not None:
+                        vector_db_params['vector_size'] = vector_size
+                vector_db_name = self._create_persistent_pgvector(vector_db_params)
+
+                # create table in vectordb before creating KB
+                self.session.datahub.get(vector_db_name).integration_handler.create_table(
+                    vector_table_name
+                )
             else:
                 # create chroma db with same name
                 vector_table_name = "default_collection"
@@ -707,10 +728,14 @@ class KnowledgeBaseController:
 
         vector_database_id = self.session.integration_controller.get(vector_db_name)['id']
 
-        # create table in vectordb
-        self.session.datahub.get(vector_db_name).integration_handler.create_table(
-            vector_table_name
-        )
+        # Store sparse vector settings in params if specified
+        if is_sparse:
+            params = params or {}
+            params['vector_config'] = {
+                'is_sparse': is_sparse
+            }
+            if vector_size is not None:
+                params['vector_config']['vector_size'] = vector_size
 
         kb = db.KnowledgeBase(
             name=name,
@@ -724,16 +749,15 @@ class KnowledgeBaseController:
         db.session.commit()
         return kb
 
-    def _create_persistent_pgvector(self):
+    def _create_persistent_pgvector(self, params=None):
         """Create default vector database for knowledge base, if not specified"""
-
         vector_store_name = "kb_pgvector_store"
 
         # check if exists
         if self.session.integration_controller.get(vector_store_name):
             return vector_store_name
 
-        self.session.integration_controller.add(vector_store_name, 'pgvector', {})
+        self.session.integration_controller.add(vector_store_name, 'pgvector', params or {})
         return vector_store_name
 
     def _create_persistent_chroma(self, kb_name, engine="chromadb"):

mindsdb/interfaces/skills/retrieval_tool.py

@@ -43,10 +43,17 @@ def build_retrieval_tool(tool: dict, pred_args: dict, skill: db.Skills):
             raise ValueError(f"Knowledge base not found: {kb_name}")
 
         kb_table = executor.session.kb_controller.get_table(kb.name, kb.project_id)
+        vector_store_config = {
+            'kb_table': kb_table
+        }
+        is_sparse = tools_config.pop('is_sparse', None)
+        vector_size = tools_config.pop('vector_size', None)
+        if is_sparse is not None:
+            vector_store_config['is_sparse'] = is_sparse
+        if vector_size is not None:
+            vector_store_config['vector_size'] = vector_size
         kb_params = {
-            'vector_store_config': {
-                'kb_table': kb_table
-            }
+            'vector_store_config': vector_store_config
         }
 
         # Get embedding model from knowledge base table

mindsdb/utilities/cache.py

@@ -71,10 +71,13 @@ _CACHE_MAX_SIZE = 500
 
 
 def dataframe_checksum(df: pd.DataFrame):
-
-    return str_checksum(str(
-        df.set_axis(range(len(df.columns)), axis=1).to_records(index=False)
-    ))
+    original_columns = df.columns
+    df.columns = list(range(len(df.columns)))
+    result = hashlib.sha256(
+        str(df.values).encode()
+    ).hexdigest()
+    df.columns = original_columns
+    return result
 
 
 def json_checksum(obj: t.Union[dict, list]):

mindsdb/utilities/context.py

@@ -52,6 +52,15 @@ class Context:
     def load(self, storage: dict) -> None:
         self._storage.set(storage)
 
+    def metadata(self, **kwargs) -> dict:
+        return {
+            'user_id': self.user_id or "",
+            'company_id': self.company_id or "",
+            'session_id': self.session_id,
+            'user_class': self.user_class,
+            **kwargs
+        }
+
 
 _context_var = ContextVar('mindsdb.context')
 context = Context(_context_var)

mindsdb/utilities/log.py

@@ -29,6 +29,23 @@ class ColorFormatter(logging.Formatter):
         return log_fmt.format(record)
 
 
+def get_console_handler_config_level() -> int:
+    console_handler_config = app_config['logging']['handlers']['console']
+    return getattr(logging, console_handler_config["level"])
+
+
+def get_file_handler_config_level() -> int:
+    file_handler_config = app_config['logging']['handlers']['file']
+    return getattr(logging, file_handler_config["level"])
+
+
+def get_mindsdb_log_level() -> int:
+    console_handler_config_level = get_console_handler_config_level()
+    file_handler_config_level = get_file_handler_config_level()
+
+    return min(console_handler_config_level, file_handler_config_level)
+
+
 def configure_logging():
     handlers_config = {}
     console_handler_config = app_config['logging']['handlers']['console']
@@ -39,6 +56,7 @@ def configure_logging():
             "formatter": "f",
             "level": console_handler_config_level
         }
+
     file_handler_config = app_config['logging']['handlers']['file']
     file_handler_config_level = getattr(logging, file_handler_config["level"])
     if file_handler_config['enabled'] is True:
@@ -51,7 +69,7 @@ def configure_logging():
             "backupCount": file_handler_config["backupCount"]
         }
 
-    mindsdb_log_level = min(console_handler_config_level, file_handler_config_level)
+    mindsdb_log_level = get_mindsdb_log_level()
 
     logging_config = dict(
         version=1,
@@ -65,7 +83,7 @@ def configure_logging():
         loggers={
             "": {  # root logger
                 "handlers": list(handlers_config.keys()),
-                "level": logging.WARNING,
+                "level": mindsdb_log_level,
             },
             "__main__": {
                 "level": mindsdb_log_level,

mindsdb/utilities/otel/__init__.py

@@ -0,0 +1,206 @@
+import os
+import typing
+
+from opentelemetry import trace  # noqa: F401
+from opentelemetry import metrics  # noqa: F401
+from opentelemetry.exporter.otlp.proto.grpc._log_exporter import OTLPLogExporter as OTLPLogExporterGRPC
+from opentelemetry.exporter.otlp.proto.http._log_exporter import OTLPLogExporter as OTLPLogExporterHTTP
+from opentelemetry.sdk._logs._internal.export import LogExporter
+from opentelemetry.exporter.otlp.proto.grpc.metric_exporter import OTLPMetricExporter as OTLPMetricExporterGRPC
+from opentelemetry.exporter.otlp.proto.http.metric_exporter import OTLPMetricExporter as OTLPMetricExporterHTTP
+from opentelemetry.sdk.metrics.export import MetricExporter, ConsoleMetricExporter
+from opentelemetry.exporter.otlp.proto.grpc.trace_exporter import OTLPSpanExporter as OTLPSpanExporterGRPC
+from opentelemetry.exporter.otlp.proto.http.trace_exporter import OTLPSpanExporter as OTLPSpanExporterHTTP
+from opentelemetry.sdk.trace.export import SpanExporter, ConsoleSpanExporter
+from opentelemetry.sdk.resources import Resource
+from opentelemetry.sdk.trace.sampling import TraceIdRatioBased
+
+from mindsdb.utilities.otel.logger import setup_logger
+from mindsdb.utilities.otel.meter import setup_meter
+from mindsdb.utilities.otel.tracer import setup_tracer
+from mindsdb.utilities.utils import parse_csv_attributes
+from mindsdb.utilities import log
+
+logger = log.getLogger(__name__)
+
+# Check OpenTelemetry exporter type
+OTEL_EXPORTER_TYPE = os.getenv("OTEL_EXPORTER_TYPE", "console")  # console or otlp
+
+# Define OpenTelemetry exporter protocol
+OTEL_EXPORTER_PROTOCOL = os.getenv("OTEL_EXPORTER_PROTOCOL", "grpc")  # grpc or http
+
+# Define OTLP endpoint. If not set, the default OTLP endpoint will be used
+OTEL_OTLP_ENDPOINT = os.getenv("OTEL_OTLP_ENDPOINT", "http://localhost:4317")
+
+# Define OTLP logging endpoint. If not set, the default OTLP logging endpoint will be used
+OTEL_OTLP_LOGGING_ENDPOINT = os.getenv("OTEL_OTLP_LOGGING_ENDPOINT", OTEL_OTLP_ENDPOINT)
+
+# Define OTLP tracing endpoint. If not set, the default OTLP tracing endpoint will be used
+OTEL_OTLP_TRACING_ENDPOINT = os.getenv("OTEL_OTLP_TRACING_ENDPOINT", OTEL_OTLP_ENDPOINT)
+
+# Define OTLP metrics endpoint. If not set, the default OTLP metrics endpoint will be used
+OTEL_OTLP_METRICS_ENDPOINT = os.getenv("OTEL_OTLP_METRICS_ENDPOINT", OTEL_OTLP_ENDPOINT)
+
+# Define service name
+OTEL_SERVICE_NAME = os.getenv("OTEL_SERVICE_NAME", "mindsdb")
+
+# Define service instace ID
+OTEL_SERVICE_INSTANCE_ID = os.getenv("OTEL_SERVICE_INSTANCE_ID", "mindsdb-instance")
+
+# The name of the environment we"re on, by default local for development, this is set differently per-env in our Helm
+# chart values files
+OTEL_SERVICE_ENVIRONMENT = os.getenv("OTEL_SERVICE_ENVIRONMENT", "local").lower()
+
+# Define service release
+OTEL_SERVICE_RELEASE = os.getenv("OTEL_SERVICE_RELEASE", "local").lower()
+
+# Define how often to capture traces
+OTEL_TRACE_SAMPLE_RATE = float(os.getenv("OTEL_TRACE_SAMPLE_RATE", "1.0"))
+
+# Define extra attributes
+OTEL_EXTRA_ATTRIBUTES = os.getenv("OTEL_EXTRA_ATTRIBUTES", "")
+
+# By default, we have Open Telemetry SDK enabled on all envs, except for local which is disabled by default.
+OTEL_SDK_DISABLED = (os.getenv("OTEL_SDK_DISABLED", "false").lower() == "true"
+                     or os.getenv("OTEL_SERVICE_ENVIRONMENT", "local").lower() == "local")
+
+# Define if OpenTelemetry logging is disabled. By default, it is disabled.
+OTEL_LOGGING_DISABLED = os.getenv("OTEL_LOGGING_DISABLED", "true").lower() == "true"
+
+# Define if OpenTelemetry tracing is disabled. By default, it is enabled.
+OTEL_TRACING_DISABLED = os.getenv("OTEL_TRACING_DISABLED", "false").lower() == "true"
+
+# Define if OpenTelemetry metrics is disabled. By default, it is disabled.
+OTEL_METRICS_DISABLED = os.getenv("OTEL_METRICS_DISABLED", "true").lower() == "true"
+
+# If you want to enable Open Telemetry on local for some reason please set OTEL_SDK_FORCE_RUN to true
+OTEL_SDK_FORCE_RUN = os.getenv("OTEL_SDK_FORCE_RUN", "false").lower() == "true"
+
+
+def get_otel_attributes() -> dict:
+    """
+    Get OpenTelemetry attributes
+
+    Returns:
+        dict: OpenTelemetry attributes
+    """
+
+    base_attributes = {
+        "service.name": OTEL_SERVICE_NAME,
+        "service.instance.id": OTEL_SERVICE_INSTANCE_ID,
+        "environment": OTEL_SERVICE_ENVIRONMENT,
+        "release": OTEL_SERVICE_RELEASE,
+    }
+
+    extra_attributes = {}
+    try:
+        extra_attributes = parse_csv_attributes(OTEL_EXTRA_ATTRIBUTES)
+    except Exception as e:
+        logger.error(f"Failed to parse OTEL_EXTRA_ATTRIBUTES: {e}")
+
+    attributes = {**extra_attributes, **base_attributes}  # Base attributes take precedence over extra attributes
+
+    return attributes
+
+
+def get_logging_exporter() -> typing.Optional[LogExporter]:
+    """
+    Get OpenTelemetry logging exporter.
+
+    Returns:
+        OTLPLogExporter: OpenTelemetry logging exporter
+    """
+
+    if OTEL_EXPORTER_TYPE == "otlp":
+
+        if OTEL_EXPORTER_PROTOCOL == "grpc":
+            return OTLPLogExporterGRPC(
+                endpoint=OTEL_OTLP_LOGGING_ENDPOINT,
+                insecure=True
+            )
+
+        elif OTEL_EXPORTER_PROTOCOL == "http":
+            return OTLPLogExporterHTTP(
+                endpoint=OTEL_OTLP_LOGGING_ENDPOINT
+            )
+
+    return None
+
+
+def get_span_exporter() -> SpanExporter:
+    """
+    Get OpenTelemetry span exporter
+
+    Returns:
+        OTLPSpanExporter: OpenTelemetry span exporter
+    """
+
+    if OTEL_EXPORTER_TYPE == "otlp":
+
+        if OTEL_EXPORTER_PROTOCOL == "grpc":
+            return OTLPSpanExporterGRPC(
+                endpoint=OTEL_OTLP_TRACING_ENDPOINT,
+                insecure=True
+            )
+
+        elif OTEL_EXPORTER_PROTOCOL == "http":
+            return OTLPSpanExporterHTTP(
+                endpoint=OTEL_OTLP_TRACING_ENDPOINT
+            )
+
+    return ConsoleSpanExporter()
+
+
+def get_metrics_exporter() -> typing.Optional[MetricExporter]:
+    """
+    Get OpenTelemetry metrics exporter
+
+    Returns:
+        OTLPLogExporter: OpenTelemetry metrics exporter
+    """
+
+    if OTEL_EXPORTER_TYPE == "otlp":
+
+        if OTEL_EXPORTER_PROTOCOL == "grpc":
+            return OTLPMetricExporterGRPC(
+                endpoint=OTEL_OTLP_METRICS_ENDPOINT,
+                insecure=True
+            )
+
+        elif OTEL_EXPORTER_PROTOCOL == "http":
+            return OTLPMetricExporterHTTP(
+                endpoint=OTEL_OTLP_METRICS_ENDPOINT
+            )
+
+    return ConsoleMetricExporter()
+
+
+if not OTEL_SDK_DISABLED or OTEL_SDK_FORCE_RUN:
+    logger.info("OpenTelemetry enabled")
+    logger.info(f"OpenTelemetry exporter type: {OTEL_EXPORTER_TYPE}")
+    logger.info(f"OpenTelemetry service name: {OTEL_SERVICE_NAME}")
+    logger.info(f"OpenTelemetry service environment: {OTEL_SERVICE_ENVIRONMENT}")
+    logger.info(f"OpenTelemetry service release: {OTEL_SERVICE_RELEASE}")
+    logger.info(f"OpenTelemetry trace sample rate: {OTEL_TRACE_SAMPLE_RATE}")
+    logger.info(f"OpenTelemetry extra attributes: {OTEL_EXTRA_ATTRIBUTES}")
+
+    # Define OpenTelemetry resources (e.g., service name)
+    attributes = get_otel_attributes()
+
+    # Define OpenTelemetry sampler
+    sampler = TraceIdRatioBased(OTEL_TRACE_SAMPLE_RATE)
+
+    # Define OpenTelemetry resources (e.g., service name)
+    resource = Resource(attributes=attributes)
+
+    if not OTEL_LOGGING_DISABLED:
+        logger.info("OpenTelemetry Logging is enabled")
+        setup_logger(resource, get_logging_exporter())
+
+    if not OTEL_TRACING_DISABLED:
+        logger.info("OpenTelemetry Tracing is enabled")
+        setup_tracer(resource, sampler, get_span_exporter())
+
+    if not OTEL_METRICS_DISABLED:
+        logger.info("OpenTelemetry Metrics is enabled")
+        setup_meter(resource, get_metrics_exporter())

mindsdb/utilities/otel/logger.py

@@ -0,0 +1,25 @@
+import logging
+
+from opentelemetry._logs import set_logger_provider
+from opentelemetry.sdk._logs._internal.export import LogExporter
+from opentelemetry.sdk._logs import LoggerProvider, LoggingHandler
+from opentelemetry.sdk._logs.export import BatchLogRecordProcessor
+from opentelemetry.sdk.resources import Resource
+
+from mindsdb.utilities.log import get_mindsdb_log_level
+
+
+def setup_logger(resource: Resource, exporter: LogExporter) -> None:
+    """
+    Setup OpenTelemetry logging
+    """
+    mindsdb_log_level = get_mindsdb_log_level()
+
+    logger_provider = LoggerProvider(resource=resource)
+    set_logger_provider(logger_provider)
+
+    logger_provider.add_log_record_processor(BatchLogRecordProcessor(exporter))
+    handler = LoggingHandler(level=mindsdb_log_level, logger_provider=logger_provider)
+
+    # Attach OTLP handler to root logger
+    logging.getLogger().addHandler(handler)

mindsdb/utilities/otel/meter.py

@@ -0,0 +1,19 @@
+from opentelemetry import metrics
+from opentelemetry.sdk.metrics import MeterProvider
+from opentelemetry.sdk.resources import Resource
+from opentelemetry.sdk.metrics.export import (
+    MetricExporter,
+    PeriodicExportingMetricReader,
+)
+
+
+def setup_meter(resource: Resource, exporter: MetricExporter) -> None:
+    """
+    Setup OpenTelemetry metrics
+    """
+
+    metric_reader = PeriodicExportingMetricReader(exporter=exporter)
+    provider = MeterProvider(resource=resource, metric_readers=[metric_reader])
+
+    # Sets the global default meter provider
+    metrics.set_meter_provider(provider)