zaturn 0.1.7__py3-none-any.whl → 0.2.0__py3-none-any.whl

zaturn/mcp/__init__.py

@@ -0,0 +1,97 @@
+import argparse
+import os
+import platformdirs
+import pkg_resources
+import sys
+
+from fastmcp import FastMCP
+
+from zaturn.tools import ZaturnTools
+
+# Basic Setup
+USER_DATA_DIR = platformdirs.user_data_dir('zaturn', 'zaturn')
+SOURCES_FILE = os.path.join(USER_DATA_DIR, 'sources.txt')
+
+# Parse command line args
+parser = argparse.ArgumentParser(
+    description="Zaturn MCP: A read-only BI tool for analyzing various data sources"
+)
+parser.add_argument('sources', nargs=argparse.REMAINDER, default=[], 
+    help='Data source (can be specified multiple times). Can be SQLite, MySQL, PostgreSQL connection string, or a path to CSV, Parquet, or DuckDB file.'
+)
+args = parser.parse_args()
+
+source_list = []
+if os.path.exists(SOURCES_FILE):
+    with open(SOURCES_FILE) as f:
+        source_list = [line.strip('\n') for line in f.readlines() if line.strip('\n')]
+
+if not source_list:
+    source_list = args.sources
+
+if not source_list:
+    source_list = [
+        pkg_resources.resource_filename(
+            'zaturn',
+            os.path.join('mcp', 'example_data', 'all_pokemon_data.csv')
+        )
+    ]
+    print("No data sources provided. Loading example dataset for demonstration.")
+    print(f"\nTo load your datasets, add them to {SOURCES_FILE} (one source URL or full file path per line)")
+    print("\nOr use command line args to specify data sources:")
+    print("zaturn_mcp sqlite:///path/to/mydata.db /path/to/my_file.csv")
+    print(f"\nNOTE: Sources in command line args will be ignored if sources are found in {SOURCES_FILE}")
+
+    
+SOURCES = {}
+for s in source_list:
+    source = s.lower()
+    if source.startswith('sqlite://'):
+        source_type = 'sqlite'
+        source_name = source.split('/')[-1].split('?')[0].split('.db')[0]
+    elif source.startswith('postgresql://'):
+        source_type = 'postgresql'
+        source_name = source.split('/')[-1].split('?')[0]
+    elif source.startswith("mysql://") or source.startswith("mysql+pymysql://"):
+        source_type = 'mysql'
+        s = s.replace('mysql://', 'mysql+pymysql://')
+        source_name = source.split('/')[-1].split('?')[0]
+    elif source.startswith('clickhouse://'):
+        source_type = 'clickhouse'
+        source_name = source.split('/')[-1].split('?')[0]
+    elif source.endswith(".duckdb"):
+        source_type = "duckdb"
+        source_name = source.split('/')[-1].split('.')[0]
+    elif source.endswith(".csv"):
+        source_type = "csv"
+        source_name = source.split('/')[-1].split('.')[0]
+    elif source.endswith(".parquet") or source.endswith(".pq"):
+        source_type = "parquet"
+        source_name = source.split('/')[-1].split('.')[0]
+    else:
+        continue
+
+    source_id = f'{source_name}-{source_type}'
+    if source_id in SOURCES:
+        i = 2
+        while True:
+            source_id = f'{source_name}{i}-{source_type}'
+            if source_id not in SOURCES:
+                break
+            i += 1
+
+    SOURCES[source_id] = {'url': s, 'source_type': source_type}
+
+
+def ZaturnMCP(sources):
+    zaturn_tools = ZaturnTools(sources)
+    zaturn_mcp = FastMCP()
+    for tool in zaturn_tools.tools:
+        zaturn_mcp.add_tool(tool)
+
+    return zaturn_mcp
+
+
+def main():
+    ZaturnMCP(SOURCES).run()
+

zaturn/studio/__init__.py

@@ -0,0 +1,5 @@
+from zaturn.studio.app import app
+
+
+def main():
+    app.run(port=6066, debug=True)

zaturn/studio/agent_wrapper.py

@@ -0,0 +1,131 @@
+import asyncio
+import json
+
+from function_schema import get_function_schema
+import httpx
+from mcp.types import ImageContent
+
+
+class Agent:
+
+    def __init__(self, 
+        endpoint: str,
+        api_key: str,
+        model: str,
+        tools: list = [],
+        image_input: bool = False,
+    ):
+    
+        self._post_url = f'{endpoint}/chat/completions'
+        self._api_key = api_key
+        self._model = model
+        self._image_input = image_input
+        self._system_message = {
+            'role': 'system',
+            'content': """
+                You are a helpful data analysis assistant.
+                Use only the tool provided data sources to process user inputs.
+                Do not use external sources or your own knowledge base.
+                Also, the tool outputs are shown to the user.
+                So, please avoid repeating the tool outputs in the generated text.
+                Use list_sources and describe_table whenever needed, 
+                do not prompt the user for source names and column names.
+            """,
+        }
+        
+        self._tools = []
+        self._tool_map = {}
+        for tool in tools:
+            tool_schema = get_function_schema(tool)
+            self._tools.append({
+                'type': 'function', 
+                'function': tool_schema,
+            })
+            self._tool_map[tool_schema['name']] = tool
+
+
+    def _prepare_input_messages(self, messages):
+        input_messages = [self._system_message]
+        for message in messages:
+            if message['role']!='tool':
+                input_messages.append(message)
+            elif type(message['content']) is not list:
+                input_messages.append(message)
+            else:
+                new_content = []
+                image_content = None
+                for content in message['content']:
+                    if content['type']=='image_url':
+                        image_content = content
+                        new_content.append({
+                            'type': 'text',
+                            'text': 'Tool call returned an image to the user.',
+                        })
+                    else:
+                        new_content.append(content)
+                input_messages.append({
+                    'role': message['role'],
+                    'tool_call_id': message['tool_call_id'],
+                    'name': message['name'],
+                    'content': new_content,
+                })
+
+        return input_messages                
+        
+    
+    def run(self, messages):
+        if type(messages) is str:
+            messages = [{'role': 'user', 'content': messages}]
+
+        while True:
+            res = httpx.post(
+                url = self._post_url,
+                headers = {
+                    'Authorization': f'Bearer {self._api_key}'
+                },
+                json = {
+                    'model': self._model,
+                    'messages': self._prepare_input_messages(messages),
+                    'tools': self._tools,
+                    'reasoning': {'exclude': True},
+                }
+            )
+
+            print(res.text)
+            resj = res.json()
+            reply = resj['choices'][0]['message']
+            messages.append(reply)
+
+            tool_calls = reply.get('tool_calls')
+            if tool_calls:
+                for tool_call in tool_calls:
+                    tool_name = tool_call['function']['name']
+                    tool_args = json.loads(tool_call['function']['arguments'])
+                    tool_response = self._tool_map[tool_name](**tool_args)
+                    if type(tool_response) is ImageContent:
+                        b64_data = tool_response.data
+                        data_url = f'data:image/png;base64,{b64_data}'
+                        content = [{
+                            'type': 'image_url',
+                            'image_url': {
+                                "url": data_url,
+                            }                            
+                        }]
+                    else:
+                        content = [{
+                            'type': 'text',
+                            'text': json.dumps(tool_response)
+                        }]
+                        
+                    messages.append({
+                        'role': 'tool',
+                        'tool_call_id': tool_call['id'],
+                        'name': tool_name,
+                        'content': content
+                    })
+            else:
+                break
+
+        return messages
+
+    

zaturn/studio/app.py

@@ -0,0 +1,288 @@
+from datetime import datetime
+import json
+
+from flask import Flask, make_response, redirect, request, render_template
+import httpx
+import mistune
+import tomli_w
+from werkzeug.utils import secure_filename
+
+from zaturn.studio import storage, agent_wrapper
+from zaturn.tools import ZaturnTools
+
+
+app = Flask(__name__)
+app.config['state'] = storage.load_state()
+
+
+def boost(content: str, fallback=None, retarget=None, reswap=None, push_url=None) -> str:
+    if request.headers.get('hx-boosted'):
+        response = make_response(content)
+        if retarget:
+            response.headers['hx-retarget'] = retarget
+        if reswap:
+            response.headers['hx-reswap'] = reswap
+        if push_url:
+            response.headers['hx-push-url'] = push_url
+        return response
+    else:
+        if fallback:
+            return fallback 
+        else:
+            slugs = storage.list_chats()
+            return render_template('_shell.html', content=content, slugs=slugs)
+
+
+@app.route('/')
+def home() -> str:
+    state = app.config['state']
+    if state.get('api_key') and state.get('sources'):
+        return boost(render_template('new_conversation.html'))
+    elif state.get('api_key'):
+        return boost(render_template('manage_sources.html'))
+    else:
+        return boost(render_template('setup_prompt.html'))
+
+
+@app.route('/settings')
+def settings() -> str:
+    return boost(render_template(
+        'settings.html', 
+        current = app.config['state'],
+        updated = request.args.get('updated'),
+    ))
+
+
+@app.route('/save_settings', methods=['POST'])
+def save_settings() -> str:
+    app.config['state']['api_key'] = request.form.get('api_key')
+
+    api_model = request.form.get('api_model').strip('/')
+    api_endpoint = request.form.get('api_endpoint').strip('/')
+    app.config['state']['api_model'] = api_model
+    app.config['state']['api_endpoint'] = api_endpoint
+    app.config['state']['api_image_input'] = False
+    
+    try:
+        model_info = httpx.get(
+            url = f'{api_endpoint}/models/{api_model}/endpoints'
+        ).json()
+        input_modalities = model_info['data']['architecture']['input_modalities']
+        if 'image' in input_modalities:
+            app.config['state']['api_image_input'] = True
+    except:
+        pass
+    storage.save_state(app.config['state'])
+    return redirect(f'/settings?updated={datetime.now().isoformat().split(".")[0]}')
+
+
+@app.route('/sources/manage')
+def manage_sources() -> str:
+    return boost(render_template(
+        'manage_sources.html',
+        sources = app.config['state'].get('sources', {})
+    ))
+
+
+@app.route('/source/toggle/', methods=['POST'])
+def source_toggle_active():
+    key = request.form['key']
+    new_active = request.form['new_status']=='active'
+    app.config['state']['sources'][key]['active'] = new_active
+    storage.save_state(app.config['state'])
+    
+    return boost(
+        render_template('c_source_card.html', key=key, active=new_active),
+        fallback = redirect('/sources/manage'),
+        retarget = f'#source-card-{key}',
+        reswap = 'outerHTML',
+        push_url = 'false',
+    )
+    
+
+@app.route('/upload_datafile', methods=['POST'])
+def upload_datafile() -> str:
+    datafile = request.files.get('datafile')
+    filename = secure_filename(datafile.filename)
+    
+    saved_path = storage.save_datafile(datafile, filename)
+    stem = saved_path.stem.replace('.', '_')
+    ext = saved_path.suffix.strip('.').lower()
+
+    app.config['state']['sources'] = app.config['state'].get('sources', {})
+    if ext in ['csv']:
+        app.config['state']['sources'][f'{stem}-csv'] = {
+            'source_type': 'csv',
+            'url': str(saved_path),
+            'active': True,
+        }
+    elif ext in ['parquet', 'pq']:
+        app.config['state']['sources'][f'{stem}-parquet'] = {
+            'source_type': 'parquet',
+            'url': str(saved_path),
+            'active': True,
+        }
+    elif ext in ['duckdb']:
+        app.config['state']['sources'][f'{stem}-duckdb'] = {
+            'source_type': 'duckdb',
+            'url': str(saved_path),
+            'active': True,
+        }
+    elif ext in ['db', 'sqlite', 'sqlite3']:
+        app.config['state']['sources'][f'{stem}-sqlite'] = {
+            'source_type': 'sqlite',
+            'url': f'sqlite:///{str(saved_path)}',
+            'active': True,
+        }
+    else:
+        storage.remove_datafile(saved_path)
+
+    storage.save_state(app.config['state'])
+    
+    return redirect('/sources/manage')
+
+
+@app.route('/add_dataurl', methods=['POST'])
+def add_dataurl():
+    url = request.form['db_url']
+    name = url.split('/')[-1].split('?')[0]
+    
+    if url.startswith("postgresql://"):
+        app.config['state']['sources'][f'{name}-postgresql'] = {
+            'source_type': 'postgresql',
+            'url': url,
+            'active': True,
+        }
+    elif url.startswith("mysql://"):
+        app.config['state']['sources'][f'{name}-mysql'] = {
+            'source_type': 'mysql',
+            'url': url,
+            'active': True,
+        }
+    elif url.startswith("clickhouse://"):
+        app.config['state']['sources'][f'{name}-clickhouse'] = {
+            'source_type': 'clickhouse',
+            'url': url,
+            'active': True,
+        }
+    else:
+        pass
+
+    storage.save_state(app.config['state'])
+    return redirect('/sources/manage')
+    
+
+@app.route('/source/delete', methods=['POST'])
+def delete_source():
+    key = request.form['key']
+    source = app.config['state']['sources'][key]
+    if source['source_type'] in ['csv', 'parquet', 'sqlite', 'duckdb']:
+        storage.remove_datafile(source['url'])
+
+    del app.config['state']['sources'][key]
+    storage.save_state(app.config['state'])
+    return redirect('/sources/manage')
+
+
+def get_active_sources():
+    sources = {}
+    for key in app.config['state']['sources']:
+        source = app.config['state']['sources'][key]
+        if source['active']:
+            sources[key] = source
+    return sources
+
+
+def prepare_chat_for_render(chat):
+    fn_calls = {}
+    for msg in chat['messages']:
+        if msg.get('role')=='assistant':
+            if msg.get('tool_calls'):
+                msg['is_tool_call'] = True
+                for tool_call in msg['tool_calls']:
+                    fn_call = tool_call['function']
+                    fn_call['arguments'] = tomli_w.dumps(
+                        json.loads(fn_call['arguments'])
+                    ).replace('\n', '<br>')
+                    fn_calls[tool_call['id']] = fn_call
+            else:
+                msg['html'] = mistune.html(msg['content'])
+        if msg.get('role')=='tool':
+            msg['call_details'] = fn_calls[msg['tool_call_id']]
+            if type(msg['content']) is str:
+                msg['html'] = mistune.html(json.loads(msg['text']))
+            elif type(msg['content']) is list:
+                msg['html'] = ''
+                for content in msg['content']:
+                    if content['type'] == 'image_url':
+                        data_url = content['image_url']['url']
+                        msg['html'] += f'<img src="{data_url}">'
+                    else:
+                        msg['html'] += mistune.html(json.loads(content['text']))
+                
+    return chat
+
+
+@app.route('/create_new_chat', methods=['POST'])
+def create_new_chat():
+    question = request.form['question']
+    slug = storage.create_chat(question)
+    chat = storage.load_chat(slug)
+
+    state = app.config['state']
+    agent = agent_wrapper.Agent(
+        endpoint = state['api_endpoint'],
+        api_key = state['api_key'],
+        model = state['api_model'],
+        tools = ZaturnTools(get_active_sources()).tools,
+        image_input = state['api_image_input'],
+    )
+    chat['messages'] = agent.run(chat['messages'])
+    storage.save_chat(slug, chat)
+    chat = prepare_chat_for_render(chat)
+    
+    return boost(
+        ''.join([
+            render_template('nav.html', slugs=storage.list_chats()),
+            '<main id="main">',
+            render_template('chat.html', chat=chat),
+            '</main>'
+        ]),
+        reswap = 'multi:#sidebar,#main',
+        push_url = f'/c/{slug}',
+        fallback = redirect(f'/c/{slug}'),
+    )
+
+
+@app.route('/c/<slug>')
+def show_chat(slug: str):
+    chat = prepare_chat_for_render(storage.load_chat(slug))
+    return boost(render_template('chat.html', chat=chat))
+
+
+@app.route('/follow_up_message', methods=['POST'])
+def follow_up_message():
+    slug = request.form['slug']
+    chat = storage.load_chat(slug)
+    chat['messages'].append({
+        'role': 'user', 
+        'content': request.form['question'],
+    })
+
+    state = app.config['state']
+    agent = agent_wrapper.Agent(
+        endpoint = state['api_endpoint'],
+        api_key = state['api_key'],
+        model = state['api_model'],
+        tools = ZaturnTools(get_active_sources()).tools,
+        image_input = state['api_image_input'],
+    )
+    chat['messages'] = agent.run(chat['messages'])
+    storage.save_chat(slug, chat)
+    chat = prepare_chat_for_render(chat)
+
+    return boost(
+        render_template('chat.html', chat=chat),
+        push_url = 'false',
+        reswap = 'innerHTML scroll:bottom',
+    )

zaturn/studio/static/js/htmx-multi-swap.js

@@ -0,0 +1,44 @@
+(function() {
+  /** @type {import("../htmx").HtmxInternalApi} */
+  var api
+
+  htmx.defineExtension('multi-swap', {
+    init: function(apiRef) {
+      api = apiRef
+    },
+    isInlineSwap: function(swapStyle) {
+      return swapStyle.indexOf('multi:') === 0
+    },
+    handleSwap: function(swapStyle, target, fragment, settleInfo) {
+      if (swapStyle.indexOf('multi:') === 0) {
+        var selectorToSwapStyle = {}
+        var elements = swapStyle.replace(/^multi\s*:\s*/, '').split(/\s*,\s*/)
+
+        elements.forEach(function(element) {
+          var split = element.split(/\s*:\s*/)
+          var elementSelector = split[0]
+          var elementSwapStyle = typeof (split[1]) !== 'undefined' ? split[1] : 'innerHTML'
+
+          if (elementSelector.charAt(0) !== '#') {
+            console.error("HTMX multi-swap: unsupported selector '" + elementSelector + "'. Only ID selectors starting with '#' are supported.")
+            return
+          }
+
+          selectorToSwapStyle[elementSelector] = elementSwapStyle
+        })
+
+        for (var selector in selectorToSwapStyle) {
+          var swapStyle = selectorToSwapStyle[selector]
+          var elementToSwap = fragment.querySelector(selector)
+          if (elementToSwap) {
+            api.oobSwap(swapStyle, elementToSwap, settleInfo)
+          } else {
+            console.warn("HTMX multi-swap: selector '" + selector + "' not found in source content.")
+          }
+        }
+
+        return true
+      }
+    }
+  })
+})()