MindsDB 1.2.9__py3-none-any.whl → 24.12.4.0__py3-none-any.whl

mindsdb/interfaces/agents/langfuse_callback_handler.py

@@ -0,0 +1,163 @@
+from typing import Any, Dict, Union, Optional, List
+from uuid import uuid4
+import datetime
+import os
+
+from langchain_core.callbacks.base import BaseCallbackHandler
+
+from mindsdb.utilities import log
+from mindsdb.interfaces.storage import db
+
+logger = log.getLogger(__name__)
+logger.setLevel('DEBUG')
+
+
+class LangfuseCallbackHandler(BaseCallbackHandler):
+    """Langchain callback handler that traces tool & chain executions using Langfuse."""
+
+    def __init__(self, langfuse, trace_id: Optional[str] = None, observation_id: Optional[str] = None):
+        self.langfuse = langfuse
+        self.chain_uuid_to_span = {}
+        self.action_uuid_to_span = {}
+        # if these are not available, we generate some UUIDs
+        self.trace_id = trace_id or uuid4().hex
+        self.observation_id = observation_id or uuid4().hex
+
+    def on_tool_start(
+            self, serialized: Dict[str, Any], input_str: str, **kwargs: Any
+    ) -> Any:
+        """Run when tool starts running."""
+        parent_run_uuid = kwargs.get('parent_run_id', uuid4()).hex
+        action_span = self.action_uuid_to_span.get(parent_run_uuid)
+        if action_span is None:
+            return
+        metadata = {
+            'tool_name': serialized.get("name", "tool"),
+            'started': datetime.datetime.now().isoformat()
+        }
+        action_span.update(metadata=metadata)
+
+    def on_tool_end(self, output: str, **kwargs: Any) -> Any:
+        """Run when tool ends running."""
+        parent_run_uuid = kwargs.get('parent_run_id', uuid4()).hex
+        action_span = self.action_uuid_to_span.get(parent_run_uuid)
+        if action_span is None:
+            return
+        action_span.update(
+            output=output,  # tool output is action output (unless superseded by a global action output)
+            metadata={'finished': datetime.datetime.now().isoformat()}
+        )
+
+    def on_tool_error(
+            self, error: Union[Exception, KeyboardInterrupt], **kwargs: Any
+    ) -> Any:
+        """Run when tool errors."""
+        parent_run_uuid = kwargs.get('parent_run_id', uuid4()).hex
+        action_span = self.action_uuid_to_span.get(parent_run_uuid)
+        if action_span is None:
+            return
+        try:
+            error_str = str(error)
+        except Exception:
+            error_str = "Couldn't get error string."
+        action_span.update(metadata={'error_description': error_str})
+
+    def on_chain_start(
+            self, serialized: Dict[str, Any], inputs: Dict[str, Any], **kwargs: Any
+    ) -> Any:
+        """Run when chain starts running."""
+        run_uuid = kwargs.get('run_id', uuid4()).hex
+        chain_span = self.langfuse.span(
+            name=f'{serialized.get("name", "chain")}-{run_uuid}',
+            trace_id=self.trace_id,
+            parent_observation_id=self.observation_id,
+            input=str(inputs)
+        )
+        self.chain_uuid_to_span[run_uuid] = chain_span
+
+    def on_chain_end(self, outputs: Dict[str, Any], **kwargs: Any) -> Any:
+        """Run when chain ends running."""
+        chain_uuid = kwargs.get('run_id', uuid4()).hex
+        if chain_uuid not in self.chain_uuid_to_span:
+            return
+        chain_span = self.chain_uuid_to_span.pop(chain_uuid)
+        if chain_span is None:
+            return
+        chain_span.update(output=str(outputs))
+        chain_span.end()
+
+    def on_chain_error(self, error: Union[Exception, KeyboardInterrupt], **kwargs: Any) -> Any:
+        """Run when chain errors."""
+        # Do nothing for now.
+        pass
+
+    def on_agent_action(self, action, **kwargs: Any) -> Any:
+        """Run on agent action."""
+        # Do nothing for now.
+        run_uuid = kwargs.get('run_id', uuid4()).hex
+        action_span = self.langfuse.span(
+            name=f'{getattr(action, "type", "action")}-{getattr(action, "tool", "")}-{run_uuid}',
+            trace_id=self.trace_id,
+            parent_observation_id=self.observation_id,
+            input=str(action)
+        )
+        self.action_uuid_to_span[run_uuid] = action_span
+
+    def on_agent_finish(self, finish, **kwargs: Any) -> Any:
+        """Run on agent end."""
+        # Do nothing for now.
+        run_uuid = kwargs.get('run_id', uuid4()).hex
+        if run_uuid not in self.action_uuid_to_span:
+            return
+        action_span = self.action_uuid_to_span.pop(run_uuid)
+        if action_span is None:
+            return
+        if finish is not None:
+            action_span.update(output=finish)  # supersedes tool output
+        action_span.end()
+
+    def auth_check(self):
+        if self.langfuse is not None:
+            return self.langfuse.auth_check()
+        return False
+
+
+def get_metadata(model_using: Dict) -> Dict:
+    """ Generates initial metadata mapping from information provided in a model's `using` clause.
+    Includes providers and model name.
+    """
+    metadata_keys = ['provider', 'model_name', 'embedding_model_provider']  # keeps keys relevant for tracing
+    trace_metadata = {}
+    for key in metadata_keys:
+        if key in model_using:
+            trace_metadata[key] = model_using.get(key)
+    return trace_metadata
+
+
+def get_skills(agent: db.Agents) -> List:
+    """ Retrieve skills from agent `skills` attribute. Specific to agent endpoints. """
+    return [rel.skill.type for rel in agent.skills_relationships]
+
+
+def get_tags(metadata: Dict) -> List:
+    """ Retrieves tags from existing langfuse metadata (built using `get_metadata` and `get_skills`), and environment variables. """
+    trace_tags = []
+    if os.getenv('FLASK_ENV'):
+        trace_tags.append(os.getenv('FLASK_ENV'))  # Fix: use something other than flask_env
+    if 'provider' in metadata:
+        trace_tags.append(metadata['provider'])
+    return trace_tags
+
+
+def get_tool_usage(trace) -> Dict:
+    """ Retrieves tool usage information from a langfuse trace.
+    Note: assumes trace marks an action with string `AgentAction` """
+    tool_usage = {}
+    steps = [s.name for s in trace.observations]
+    for step in steps:
+        if 'AgentAction' in step:
+            tool_name = step.split('-')[1]
+            if tool_name not in tool_usage:
+                tool_usage[tool_name] = 0
+            tool_usage[tool_name] += 1
+    return tool_usage

mindsdb/interfaces/agents/mindsdb_chat_model.py

@@ -0,0 +1,196 @@
+from __future__ import annotations
+
+import logging
+from typing import (
+    Any,
+    Dict,
+    List,
+    Mapping,
+    Optional,
+)
+
+import pandas as pd
+from langchain_core.callbacks import (
+    CallbackManagerForLLMRun,
+)
+from langchain_core.language_models.chat_models import (
+    BaseChatModel,
+)
+from langchain_core.messages import (
+    AIMessage,
+    BaseMessage,
+    ChatMessage,
+    FunctionMessage,
+    HumanMessage,
+    SystemMessage,
+)
+from langchain_core.outputs import (
+    ChatGeneration,
+    ChatResult,
+)
+from pydantic import model_validator
+
+from mindsdb.interfaces.agents.constants import USER_COLUMN
+
+logger = logging.getLogger(__name__)
+
+
+def _convert_message_to_dict(message: BaseMessage) -> dict:
+    if isinstance(message, ChatMessage):
+        message_dict = {"role": message.role, "content": message.content}
+    elif isinstance(message, HumanMessage):
+        message_dict = {"role": "user", "content": message.content}
+    elif isinstance(message, AIMessage):
+        message_dict = {"role": "assistant", "content": message.content}
+        if "function_call" in message.additional_kwargs:
+            message_dict["function_call"] = message.additional_kwargs["function_call"]
+    elif isinstance(message, SystemMessage):
+        message_dict = {"role": "system", "content": message.content}
+    elif isinstance(message, FunctionMessage):
+        message_dict = {
+            "role": "function",
+            "content": message.content,
+            "name": message.name,
+        }
+    else:
+        raise ValueError(f"Got unknown type {message}")
+    if "name" in message.additional_kwargs:
+        message_dict["name"] = message.additional_kwargs["name"]
+    return message_dict
+
+
+class ChatMindsdb(BaseChatModel):
+    """A chat model that uses the Mindsdb"""
+
+    model_name: str
+    project_name: Optional[str] = 'mindsdb'
+    model_info: Optional[dict] = None
+    project_datanode: Optional[Any] = None
+
+    class Config:
+        """Configuration for this pydantic object."""
+        arbitrary_types_allowed = True
+        allow_reuse = True
+
+    @property
+    def _default_params(self) -> Dict[str, Any]:
+        return {}
+
+    def completion(
+            self, messages: List[dict]
+    ) -> Any:
+        problem_definition = self.model_info['problem_definition'].get('using', {})
+        output_col = self.model_info['predict']
+
+        # TODO create table for conversational model?
+        if len(messages) > 1:
+            content = '\n'.join([
+                f"{m['role']}: {m['content']}"
+                for m in messages
+            ])
+        else:
+            content = messages[0]['content']
+
+        record = {}
+        params = {}
+        # Default to conversational if not set.
+        mode = problem_definition.get('mode', 'conversational')
+        if mode == 'conversational' or mode == 'retrieval':
+            # flag for langchain to prevent calling agent inside of agent
+            if self.model_info['engine'] == 'langchain':
+                params['mode'] = 'chat_model'
+
+            user_column = problem_definition.get('user_column', USER_COLUMN)
+            record[user_column] = content
+
+        elif 'column' in problem_definition:
+            # input defined as 'column' param
+            record[problem_definition['column']] = content
+
+        else:
+            # failback, maybe handler supports template injection
+            params['prompt_template'] = content
+
+        predictions = self.project_datanode.predict(
+            model_name=self.model_name,
+            df=pd.DataFrame([record]),
+            params=params,
+        )
+
+        col = output_col
+        if col not in predictions.columns:
+            # get first column
+            col = predictions.columns[0]
+
+        # get first row
+        result = predictions[col][0]
+
+        # TODO token calculation
+        return {
+            'messages': [result]
+        }
+
+    @model_validator(mode='before')
+    def validate_environment(cls, values: Dict) -> Dict:
+
+        model_name = values['model_name']
+        project_name = values['project_name']
+
+        from mindsdb.api.executor.controllers import SessionController
+
+        session = SessionController()
+        session.database = 'mindsdb'
+
+        values['model_info'] = session.model_controller.get_model(model_name, project_name=project_name)
+
+        project_datanode = session.datahub.get(values['project_name'])
+
+        values["project_datanode"] = project_datanode
+
+        return values
+
+    def _generate(
+            self,
+            messages: List[BaseMessage],
+            stop: Optional[List[str]] = None,
+            run_manager: Optional[CallbackManagerForLLMRun] = None,
+            stream: Optional[bool] = None,
+            **kwargs: Any,
+    ) -> ChatResult:
+
+        message_dicts = [_convert_message_to_dict(m) for m in messages]
+
+        response = self.completion(
+            messages=message_dicts
+        )
+        return self._create_chat_result(response)
+
+    def _create_chat_result(self, response: Mapping[str, Any]) -> ChatResult:
+        generations = []
+        for content in response["messages"]:
+            message = AIMessage(content=content)
+            gen = ChatGeneration(
+                message=message,
+                generation_info=dict(finish_reason=None),
+            )
+            generations.append(gen)
+        token_usage = response.get("usage", {})
+        set_model_value = self.model_name
+        if self.model_name is not None:
+            set_model_value = self.model_name
+        llm_output = {"token_usage": token_usage, "model": set_model_value}
+        return ChatResult(generations=generations, llm_output=llm_output)
+
+    @property
+    def _identifying_params(self) -> Dict[str, Any]:
+        """Get the identifying parameters."""
+        set_model_value = self.model_name
+        if self.model_name is not None:
+            set_model_value = self.model_name
+        return {
+            "model_name": set_model_value,
+        }
+
+    @property
+    def _llm_type(self) -> str:
+        return "mindsdb"

mindsdb/interfaces/agents/mindsdb_database_agent.py

@@ -0,0 +1,54 @@
+"""
+    Wrapper around MindsDB's executor and integration controller following the implementation of the original
+    langchain.sql_database.SQLDatabase class to partly replicate its behavior.
+"""
+from typing import Any, Iterable, List, Optional
+
+from mindsdb.interfaces.skills.skill_tool import skill_tool
+from mindsdb.utilities import log
+from langchain_community.utilities import SQLDatabase
+
+logger = log.getLogger(__name__)
+
+
+class MindsDBSQL(SQLDatabase):
+    """ Can't modify signature, as LangChain does a Pydantic check."""
+    def __init__(
+        self,
+        engine=None,
+        database: Optional[str] = 'mindsdb',
+        metadata: Optional[Any] = None,
+        ignore_tables: Optional[List[str]] = None,
+        include_tables: Optional[List[str]] = None,
+        sample_rows_in_table_info: int = 3,
+        schema: Optional[str] = None,
+        indexes_in_table_info: bool = False,
+        custom_table_info: Optional[dict] = None,
+        view_support: Optional[bool] = True,
+    ):
+        # Some args above are not used in this class, but are kept for compatibility
+
+        self._sql_agent = skill_tool.get_sql_agent(
+            database,
+            include_tables,
+            ignore_tables,
+            sample_rows_in_table_info
+        )
+
+    @property
+    def dialect(self) -> str:
+        return 'mindsdb'
+
+    @property
+    def table_info(self) -> str:
+        """Information about all tables in the database."""
+        return self._sql_agent.get_table_info()
+
+    def get_usable_table_names(self) -> Iterable[str]:
+        return self._sql_agent.get_usable_table_names()
+
+    def get_table_info_no_throw(self, table_names: Optional[List[str]] = None) -> str:
+        return self._sql_agent.get_table_info_safe(table_names)
+
+    def run_no_throw(self, command: str, fetch: str = "all") -> str:
+        return self._sql_agent.query_safe(command)

mindsdb/interfaces/agents/safe_output_parser.py

@@ -0,0 +1,51 @@
+import re
+from typing import Union
+
+from langchain_core.agents import AgentAction, AgentFinish
+
+from langchain.agents.agent import AgentOutputParser
+from langchain.agents.conversational.prompt import FORMAT_INSTRUCTIONS
+
+from mindsdb.utilities import log
+
+logger = log.getLogger(__name__)
+
+
+class SafeOutputParser(AgentOutputParser):
+    '''Output parser for the conversational agent that does not throw OutputParserException.'''
+
+    ai_prefix: str = 'AI'
+    '''Prefix to use before AI output.'''
+
+    format_instructions: str = FORMAT_INSTRUCTIONS
+    '''Default formatting instructions'''
+
+    def get_format_instructions(self) -> str:
+        '''Returns formatting instructions for the given output parser.'''
+        return self.format_instructions
+
+    def parse(self, text: str) -> Union[AgentAction, AgentFinish]:
+        '''Parses outputted text from an LLM.
+
+        Args:
+            text (str): Outputted text to parse.
+
+        Returns:
+            output (str): Parsed text to an Agent step.
+        '''
+        regex = r'Action: (.*?)[\n]*Action Input:([\s\S]*)'
+        match = re.search(regex, text, re.DOTALL)
+        if match is not None:
+            action = match.group(1)
+            action_input = match.group(2)
+            return AgentAction(action.strip(), action_input.strip(' ').strip('"'), text)
+        output = text
+        if f'{self.ai_prefix}:' in text:
+            output = text.split(f'{self.ai_prefix}:')[-1].strip()
+        return AgentFinish(
+            {'output': output}, text
+        )
+
+    @property
+    def _type(self) -> str:
+        return 'conversational'

mindsdb/interfaces/chatbot/__init__.py

@@ -0,0 +1 @@
+