studio/www/html/_plugins/agent.py

from langchain_ollama import ChatOllama
from langchain_openai import AzureChatOpenAI, ChatOpenAI
import requests
from typing import Optional
import json
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.tools import StructuredTool
from langchain.agents import create_agent
from pydantic import create_model

from pydantic import BaseModel, Field, field_validator

from typing import Dict, Any, List
from langchain.agents.structured_output import ToolStrategy
import transport
import cms
import pandas as pd
import plugin_ix as plug
import os

_map ={
    'postgresql':"Select table_schema||'.'||table_name as table_name, table_schema as schema_name, column_name, data_type as type  from information_schema.columns",
    'duckdb':"SELECT table_schema||'.'||table_name as table_name , table_schema AS schema_name, column_name, data_type FROM information_schema.columns",
    'sqlite3':'SELECT * FROM sqlite_master',
    'mysql':'SELECT table_name, table_schema AS schema_name, column_name, data_type FROM information_schema.columns',
    'bigquery':'SELECT table_name, table_schema AS schema_name, column_name, data_type FROM :dataset.INFORMATION_SCHEMA.COLUMNS',
    'iceberg':'SHOW TABLES', 'mariadb':'SHOW TABLES'
}
class AnalysisItem(BaseModel):
    title: str = Field(description="Short title for this piece of analysis")
    explanation: str = Field(description="Explanation of the insight")
    # Using a structured dict or strict List[Dict]
    data: List[Dict[str, Any]] = Field(description="Relevant rows supporting this analysis point")

    @field_validator('data', mode='before')
    @classmethod
    def ensure_list_of_dicts(cls, v):
        if isinstance(v, dict):
            return [v]  # Wrap single dicts if the LLM forgot the list wrapper
        return v

class QueryOutput(BaseModel):
    query: str = Field(description="The query that was generated and executed")
    data: List[Dict[str, Any]] = Field(description="Rows returned from executing the query")
    analysis: List[AnalysisItem] = Field(description="Analysis points derived from the data")

    @field_validator('data', mode='before')
    @classmethod
    def ensure_list_of_dicts(cls, v):
        if isinstance(v, dict):
            return [v]
        return v
    
class Tool:
    def __init__(self,_dbr) :
        self._dbr = _dbr
        self._provider = _dbr._agent._engine.name


    def catalog(self):
        """
        This function returns the catalog of the current database
        """
        global _map
        _query = _map.get(self._provider)
        
        return self.read(_query)
    def schema(self,table:str) ->dict:
        """
        This function returns the schema for a given database, if table is not provided the entire catalog will be returned with the list of tables
        """
        _query = f""" {_map[self._provider]} WHERE  name = '{table}' """
        if self._provider in ['sqlite','sqlite3'] :
            _info = self.read(f"select * from {table} limit 20").dtypes
            return pd.DataFrame({"column_name":_info.dtypes.index.tolist()[1:],"data_type": [x.name for x in _info.dtypes.tolist()[1:]]}).to_dict(orient='records')
            pass
        else:
            return self.read(_query)
    def read(self,query:str)->dict:
        """
        This function will execute a query against a given label/database. The query must adhere to the schema of the tables involved.
        """
        # print (query)
        return self._dbr.read(sql=query).to_dict(orient='records')

class TransportAgent:
    """
    This class is intended to allow an LLM to interact with data-transport MCP server.
    """
    def __init__(self, **kwargs):
        """
        :backend    ollama or openai (for now)
        :kwargs     parameters that will be used for the openai,ollama classes
        """
        self.backend = kwargs.get('backend', 'ollama').lower()
        self.kwargs = kwargs.get('kwargs', {}).copy()
        # self.mcp_url = kwargs.get('mcp_url', 'http://localhost:8090')
        self.chatbot = self.instance(self.backend, **self.kwargs)
        # self.tools = self._init_mcp_tools()
        # self.agent = self._init_agent()
        registry = plug.Registry()
        if registry.has(kwargs['parser']) :
            self._parser = registry.get(kwargs['parser'])
        else:
            self._parser = None

    def instance(self, backend, **kwargs):
        """
        Returns a chat model instance for the given backend.
        Supported: ollama, azureopenai, openai.
        """
        if backend == 'azureopenai':
            return AzureChatOpenAI(**kwargs)
        elif backend == 'openai':
            return ChatOpenAI(**kwargs)
        else:
            return ChatOllama(**kwargs)
    def get(self,label,query):
        dbr = transport.get.reader(label=label)
        tools = Tool(dbr)
        # self.chatbot.bind_tools ([tools.catalog,tools.read])

        _message = f"""
     You are a specialized Data-Transport Agent designed to facilitate data access and analysis from registered        
     labels/databases. The user is working in {label} database ({dbr._agent._engine.name} database).
                                                                                                                       
     Your primary goal is to execute a requested operation based on the user's query. You MUST strictly adhere to      
     the following process and output format.  When users ask for a list of something, they just mean the records.

     **Process Flow:**   
    if and ONLY if the user's question is a database query execute the "read" function associated with the tools.
    Otherwise perform the following:
    1.  **Catalog Retrieval:** First, return the list of tables with a field describing the table. 
    schema) from the catalog.                                                                                         
    2.  **Query Generation:** Construct an accurate SQL query that directly addresses the user's request using the    
    the information pulled from the catalog. NEVER have a query project a complex object (List; JSON; Blob)
    3. Make sure the generated query adheres to constraints of {dbr._agent._engine.name} database
    4.  **Execution:** Execute the SQL query against by calling the "read" function of the tools.                                                        
    5.  **Analysis is Optional and should be provided ONLY if the user requested it:** If the user asks for an analysis, perform the necessary data processing and derive   
    meaningful insights based on the query data.                                                                    
                                                                                                                       
     **Output Format Enforcement:**                                                                                    
     You MUST respond ONLY with a single JSON object. Do not include any introductory text, explanations, or           
     conversational language outside of the JSON structure.                                                            
                                                                                                                       
     **JSON Schema Requirements:**                                                                                     
     The output MUST strictly conform to the following JSON structure, utilizing the `QueryOutput` schema defined by   
     your execution context:                                                                                           
     - **"query"**: The exact SQL query that was executed.                                                             
     - **"data"**: A list of dictionaries containing the raw data rows returned from the query execution.            
     - **"analysis"**: An optional list of analysis points. If analysis was performed, this list MUST contain          
     objects with the following structure:                                                                             
         - "title": A concise title for the derived insight. It can ONLY be 30 characters at most.                                                           
         - "data": The specific data rows supporting the analysis.                                                   
         - "text": A detailed explanation of the insight derived from the data.                                        
                                                                                                                       
     Ensure all generated data and analysis are directly traceable back to the executed query and the retrieved        
     table schemas.
        """
        _caller = create_agent(self.chatbot,[tools.catalog,tools.schema,tools.read],system_prompt=_message,response_format=QueryOutput)
        _output =  _caller.invoke({"messages":[("user",query)]})["messages"][-1].content
        if self._parser :
            return self._parser(_output)
        else:
            return _output
LLM_ARGS = {}
def InitLLMArgs (_config):
    global LLM_ARGS

    path = _config['system']['source']['ai']['config-path']
    f = open(path)
    LLM_ARGS = json.loads(f.read())
    f.close()
    #
    # We also have to initialize the plugin path
    os.environ['REGISTRY_FOLDER'] = _config['system']['source']['ai']['plugin-path']
    return LLM_ARGS

@cms.Plugin(mimetype="application/json",method="POST")
def apply(**_args):
    global LLM_ARGS
    _request = _args['request']
    _config = _args['config']
    #
    # let's make sure the _label is whitelisted
    _label = _request.json['label']
    _query = _request.json['query']
    if not LLM_ARGS :
        LLM_ARGS = InitLLMArgs(_args['config'])
    if _query.strip().lower().startswith("with") or _query.strip().lower().startswith("select") :
        dbr = transport.get.reader(label=_label)
        
        return { "data":dbr.read(sql=_query).fillna(0).to_dict(orient='records') ,"analysis":[]}
    elif _query.strip().lower().startswith("{") :
        dbr = transport.get.reader(label=_label)
        return { "data":dbr.read(json.loads(_query)).to_dict(orient='records'),"analysis":[] }

    _agent = TransportAgent(**LLM_ARGS)
    
    # we need an output parser for the models (should be loaded as a plugin)
    print ([LLM_ARGS['backend'], LLM_ARGS['kwargs']['model']])
    return _agent.get(_label,_query)
    
    # _data =  {"query": "SELECT condition_concept_id, COUNT(*) AS condition_count FROM tabular.condition_occurrence GROUP BY condition_concept_id ORDER BY condition_count DESC LIMIT 10;", "data": [{"condition_concept_id": "316866", "condition_count": 5635}, {"condition_concept_id": "4159131", "condition_count": 5462}, {"condition_concept_id": "2005200547", "condition_count": 5375}, {"condition_concept_id": "4291025", "condition_count": 5037}, {"condition_concept_id": "201826", "condition_count": 4026}, {"condition_concept_id": "433736", "condition_count": 4011}, {"condition_concept_id": "4036620", "condition_count": 3734}, {"condition_concept_id": "4317977", "condition_count": 3341}, {"condition_concept_id": "37018196", "condition_count": 2627}, {"condition_concept_id": "81902", "condition_count": 2458}], "analysis": [{"title": "Top 10 Most Frequent Conditions", "data": [{"condition_concept_id": "316866", "condition_count": 5635}, {"condition_concept_id": "4159131", "condition_count": 5462}, {"condition_concept_id": "2005200547", "condition_count": 5375}, {"condition_concept_id": "4291025", "condition_count": 5037}, {"condition_concept_id": "201826", "condition_count": 4026}, {"condition_concept_id": "433736", "condition_count": 4011}, {"condition_concept_id": "4036620", "condition_count": 3734}, {"condition_concept_id": "4317977", "condition_count": 3341}, {"condition_concept_id": "37018196", "condition_count": 2627}, {"condition_concept_id": "81902", "condition_count": 2458}], "text": "The distribution analysis of conditions, based on the top 10 most frequent condition concept IDs, shows that '316866' is the most commonly recorded condition, with 5,635 occurrences. The top four conditions (316866, 4159131, 2005200547, and 4291025) account for a significant portion of the recorded conditions, indicating these are the most prevalent diagnoses in the dataset."}]}
    # _data['data'] = pd.DataFrame(_data['data']).to_dict(orient='split',index=False)
    # for _row in _data['analysis'] :
    #     if 'data' in _row :
    #         _row['data'] = pd.DataFrame(_row['data']).to_dict(orient='records')
    # return _data
    

@cms.Plugin(mimetype="text/plain")
def enabled(**_args):
    _config = _args['config']
    return str(int('ai' in _config['system']['source'] ))