zaturn 0.2.1__py3-none-any.whl → 0.3.0__py3-none-any.whl

zaturn/mcp/__init__.py

@@ -1,10 +1,11 @@
 import argparse
+import importlib.resources
 import os
 import platformdirs
-import pkg_resources
 import sys
 
 from fastmcp import FastMCP
+from fastmcp.tools.tool import Tool
 
 from zaturn.tools import ZaturnTools
 
@@ -30,19 +31,17 @@ 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')
-        )
-    ]
+    with importlib.resources.path(
+        'zaturn.tools.example_data', 'all_pokemon_data.csv'
+    ) as source_path:
+        source_list = [str(source_path)]
+    
     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()
@@ -90,8 +89,8 @@ for s in source_list:
 def ZaturnMCP(sources):
     zaturn_tools = ZaturnTools(sources)
     zaturn_mcp = FastMCP()
-    for tool in zaturn_tools.tools:
-        zaturn_mcp.add_tool(tool)
+    for tool_function in zaturn_tools.tools:
+        zaturn_mcp.add_tool(Tool.from_function(tool_function))
 
     return zaturn_mcp
 

zaturn/studio/agent_wrapper.py

@@ -4,128 +4,52 @@ import json
 from function_schema import get_function_schema
 import httpx
 from mcp.types import ImageContent
+from pydantic_ai import Agent
+from pydantic_ai.messages import ModelMessagesTypeAdapter
+from pydantic_ai.models.openai import OpenAIChatModel, OpenAIChatModelSettings
+from pydantic_ai.providers.openai import OpenAIProvider
+from pydantic_core import to_jsonable_python
 
 
-class Agent:
+class ZaturnAgent:
 
     def __init__(self, 
         endpoint: str,
         api_key: str,
-        model: str,
+        model_name: str,
         tools: list = [],
         image_input: bool = False,
+        reasoning_effort: str = 'none',
+        system_prompt: str = '',
     ):
     
-        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},
-                }
+        model_settings = None
+        reasoning_effort = reasoning_effort.lower()
+        if reasoning_effort in ['low', 'medium', 'high']:
+            model_settings = OpenAIChatModelSettings(
+                openai_reasoning_effort = reasoning_effort,
             )
+            
+        self._agent = Agent(
+            OpenAIChatModel(
+                model_name,
+                provider = OpenAIProvider(
+                    api_key = api_key,
+                    base_url = endpoint,
+                )
+            ),
+            model_settings = model_settings,
+            system_prompt = system_prompt or None,
+            tools = tools,
+        )
+        
 
-            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
-
+    def run(self, prompt, message_history = None):
+        if message_history:
+            message_history_obj = ModelMessagesTypeAdapter.validate_python(message_history)
+            result = self._agent.run_sync(prompt, message_history = message_history_obj)
+            return to_jsonable_python(result.all_messages())
+        else:
+            result = self._agent.run_sync(prompt)
+            return to_jsonable_python(result.all_messages())
     

zaturn/studio/app.py

@@ -62,6 +62,8 @@ def save_settings() -> str:
     app.config['state']['api_model'] = api_model
     app.config['state']['api_endpoint'] = api_endpoint
     app.config['state']['api_image_input'] = False
+    app.config['state']['reasoning_effort'] = request.form.get('reasoning_effort', 'none')
+    app.config['state']['system_prompt'] = request.form.get('system_prompt').strip('\n')
     
     try:
         model_info = httpx.get(
@@ -195,31 +197,34 @@ def get_active_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']))
-                
+    
+    for message in chat['messages']:
+        for part in message['parts']:
+            if part['part_kind']=='text' and message['kind']=='response':
+                part['html_content'] = mistune.html(part['content'])
+            elif part['part_kind']=='tool-call':
+                fn_calls[part['tool_call_id']] = part
+                fn_calls[part['tool_call_id']]['timestamp'] = message['timestamp']
+            elif part['part_kind']=='tool-return':
+                fn_call = fn_calls[part['tool_call_id']]
+                part['call_details'] = {}
+                part['call_details']['name'] = fn_call['tool_name']
+
+                t1 = datetime.fromisoformat(fn_call['timestamp'])
+                t2 = datetime.fromisoformat(part['timestamp'])
+                part['call_details']['exec_time'] = (t2 - t1).seconds
+
+                part['call_details']['args_html'] = tomli_w.dumps(
+                    json.loads(fn_call['args'])
+                ).replace('\n', '<br>')
+
+                if type(part['content']) is str:
+                    part['html_content'] = mistune.html(part['content'])
+                elif type(part['content']) is dict and part['content']['type']=='image':
+                    data_url = f"data:{part['content']['mimeType']};base64,{part['content']['data']}"
+                    part['html_content'] = f'<img src="{data_url}">'
+                    
+
     return chat
 
 
@@ -230,22 +235,23 @@ def create_new_chat():
     chat = storage.load_chat(slug)
 
     state = app.config['state']
-    agent = agent_wrapper.Agent(
+    agent = agent_wrapper.ZaturnAgent(
         endpoint = state['api_endpoint'],
         api_key = state['api_key'],
-        model = state['api_model'],
+        model_name = state['api_model'],
         tools = ZaturnTools(get_active_sources()).tools,
         image_input = state['api_image_input'],
+        reasoning_effort = state['reasoning_effort'],
+        system_prompt = state['system_prompt'],
     )
-    chat['messages'] = agent.run(chat['messages'])
+    chat['messages'] = agent.run(question)
     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),
+            render_template('chat.html', chat=prepare_chat_for_render(chat)),
             '</main>'
         ]),
         reswap = 'multi:#sidebar,#main',
@@ -264,25 +270,26 @@ def show_chat(slug: str):
 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(
+    agent = agent_wrapper.ZaturnAgent(
         endpoint = state['api_endpoint'],
         api_key = state['api_key'],
-        model = state['api_model'],
+        model_name = state['api_model'],
         tools = ZaturnTools(get_active_sources()).tools,
         image_input = state['api_image_input'],
+        reasoning_effort = state['reasoning_effort'],
+        system_prompt = state['system_prompt'],
+    )
+    
+    chat['messages'] = agent.run(
+        prompt = request.form['question'],
+        message_history = chat['messages'],
     )
-    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),
+        render_template('chat.html', chat=prepare_chat_for_render(chat)),
         push_url = 'false',
         reswap = 'innerHTML scroll:bottom',
     )

zaturn/studio/storage.py

@@ -14,12 +14,24 @@ STATE_FILE = USER_DATA_DIR / 'studio.json'
 CHATS_DIR = USER_DATA_DIR / 'chats'
 os.makedirs(CHATS_DIR, exist_ok=True)
 
+DEFAULT_PROMPT = """
+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_data_sources and describe_table whenever needed, 
+do not prompt the user for source names and column names.
+""".strip('\n')
+
 
 def load_state() -> dict:
     if os.path.exists(STATE_FILE):
         with open(STATE_FILE) as f:
             state = json.loads(f.read())
             state['sources'] = state.get('sources', {})
+            state['system_prompt'] = state.get('system_prompt', DEFAULT_PROMPT)
+            state['reasoning_effort'] = state.get('reasoning_effort', 'none') 
             return state
     else:
         return {}
@@ -51,10 +63,8 @@ def create_chat(question: str):
     
     chat = {
         'slug': slug,
-        'messages': [{
-            'role': 'user',
-            'content': question,
-        }]
+        'messages': [],
+        'schema_version': 1
     }
     
     filename = CHATS_DIR / f'{slug}.json'

zaturn/studio/templates/_shell.html

@@ -19,13 +19,13 @@
         <img src="/static/logo.svg" class="logo">
       </a>
       <a href="/">
-        💬
+        {% include('icons/chat-bubble.svg') %}
       </a>
       <a href="/sources/manage" title="Manage Sources">
-        📎
+        {% include('icons/database.svg') %}
       </a>
       <a href="/settings" title="Settings">
-        🛠
+        {% include('icons/settings.svg') %}
       </a>
     </header>
     {% with slugs=slugs %}

zaturn/studio/templates/ai_message.html

@@ -1,4 +1,3 @@
 <div class="ai-message">
-  <!--<p><b class="sender">Zaturn:</b></p>-->
-  <div>{{msg['html'] | safe}}</div>
+  <div>{{part['html_content'] | safe}}</div>
 </div>

zaturn/studio/templates/c_source_card.html

@@ -5,15 +5,15 @@
     <input type="hidden" name="key" value="{{key}}">
     {% if active %}
     <input type="hidden" name="new_status" value="inactive">
-    <button class="active" title="Currently visible to LLM, Click to Toggle">Active ✅</button> 
+    <button class="active" title="Currently visible to LLM, Click to Toggle">Active {% include('icons/check-circle-solid.svg') %}</button> 
     {% else %}
     <input type="hidden" name="new_status" value="active">
-    <button class="inactive" title="Currently invisible to LLM, Click to Toggle">Inactive 😴</button>
+    <button class="inactive" title="Currently invisible to LLM, Click to Toggle">Inactive {% include('icons/warning-triangle.svg') %}</button>
     {% endif %}
   </form>
 
   <form action="/source/delete" method="POST">
     <input type="hidden" name="key" value="{{key}}">
-    <button class="danger">Delete 🗑</button>
+    <button class="danger">Delete {% include('icons/trash.svg') %}</button>
   </form>
 </div>

zaturn/studio/templates/chat.html

@@ -1,22 +1,30 @@
 <section id="chat">
-  {% for msg in chat['messages'] %}
-    {% if msg['role']=='user' %}
-      {% include('user_message.html') %}
-    {% elif msg['role']=='assistant' and not msg['is_tool_call'] %}
-      {% include('ai_message.html') %}
-    {% elif msg['role']=='tool' %}
-      {% include('function_call.html') %}
-    {% endif %}
-  {% endfor %}
+  {% if chat.get('schema_version', 0) >= 1 %}
+    {% for message in chat['messages'] %}
+      {% for part in message['parts'] %}
+        {% if part['part_kind']=='system-prompt' %}
+          {% include('chat_metadata.html') %}
+        {% elif part['part_kind']=='user-prompt' %}
+          {% include('user_message.html') %}
+        {% elif part['part_kind']=='text' and message['kind']=='response' %}
+          {% include('ai_message.html') %}
+        {% elif part['part_kind']=='tool-return' %}
+          {% include('function_call.html') %}
+        {% endif %}
+      {% endfor %}
+    {% endfor %}
 
-  <form action="/follow_up_message" method="POST">
-    <input type="hidden" name="slug" value="{{chat['slug']}}">
-    <textarea 
-      required 
-      name="question" 
-      placeholder="Type a follow up question here."
-    ></textarea>
-    <button>➡</button>
-  </form>
-  {% include('loader.html') %}
+    <form action="/follow_up_message" method="POST">
+      <input type="hidden" name="slug" value="{{chat['slug']}}">
+      <textarea 
+        required 
+        name="question" 
+        placeholder="Type a follow up question here."
+      ></textarea>
+      <button>{% include('icons/arrow-right.svg') %}</button>
+    </form>
+    {% include('loader.html') %}
+  {% else %}
+    <pre><code>{{chat | tojson(indent=2)}}</code></pre>
+  {% endif %}
 </section>

zaturn/studio/templates/chat_metadata.html

@@ -0,0 +1,10 @@
+<div class="chat-metadata">
+  <dl>
+    <dt>Model:</dt>
+    <dd>{{chat['messages'][-1]['model_name']}}</dd>
+  </dl>
+  <details class="system-prompt-message">
+    <summary>System Prompt</summary>
+    <p>{{part['content']}}</p>
+  </details>
+</div>