mulebot 0.2.0__tar.gz → 0.3.0__tar.gz

{mulebot-0.2.0 → mulebot-0.3.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: mulebot
-Version: 0.2.0
+Version: 0.3.0
 Summary: A simple chatbot for Datamule's APIs
 Home-page: https://github.com/john-friedman/mulebot
 Author: John Friedman

mulebot-0.3.0/mulebot/__init__.py

@@ -0,0 +1,5 @@
+from .agent import Agent
+from .tool import Tool, tool
+from .providers import DEFAULT_MODELS
+
+__all__ = ["Agent", "Tool", "tool", "DEFAULT_MODELS"]

{mulebot-0.2.0 → mulebot-0.3.0}/mulebot/agent.py

@@ -7,9 +7,10 @@ from time import perf_counter, time
 
 from .tools import Tools, documentation_context
 from .tool import validate_tools
+from .providers import DEFAULT_MODELS, build_request, credentials, normalize_response, resolve_model
 
 
-DEFAULT_MODEL = "deepseek/deepseek-v4.1-flash-1017:nitro"
+DEFAULT_MODEL = DEFAULT_MODELS["openrouter"]
 TOOL_LIMIT = 50
 
 
@@ -63,15 +64,13 @@ def reasoning_text(response):
 
 
 class Agent:
-    def __init__(self, api_key=None, model=None, model_api_key=None, allowed_tables=None):
-        self.model = model or os.environ.get("MULEBOT_MODEL", DEFAULT_MODEL)
-        self.model_api_key = model_api_key or os.environ.get("OPENROUTER_API_KEY")
-        api_key = api_key or os.environ.get("DATAMULE_API_KEY")
-        if not api_key or not self.model_api_key:
-            raise ValueError("Set DATAMULE_API_KEY and OPENROUTER_API_KEY.")
-        self._tool_state = Tools(api_key, allowed_tables=allowed_tables)
+    def __init__(self, provider, *, model=None, allowed_tables=None):
+        self.provider = provider
+        self.model = resolve_model(provider, model)
+        self.model_api_key = credentials(provider)
+        self._tool_state = Tools(allowed_tables=allowed_tables)
         self.instructions = INSTRUCTIONS
-        self.doc = documentation_context()
+        self.documentation = documentation_context()
         self.messages = []
         self._update_system_message()
 
@@ -100,17 +99,17 @@ class Agent:
             raise TypeError("Instructions must be a string.")
         return self.set_instructions("\n\n".join(part for part in (self.instructions, text) if part))
 
-    def set_doc(self, text):
+    def set_documentation(self, text):
         if not isinstance(text, str):
             raise TypeError("Documentation must be a string.")
-        self.doc = text
+        self.documentation = text
         self._update_system_message()
         return self
 
-    def append_doc(self, text):
+    def append_documentation(self, text):
         if not isinstance(text, str):
             raise TypeError("Documentation must be a string.")
-        return self.set_doc("\n\n".join(part for part in (self.doc, text) if part))
+        return self.set_documentation("\n\n".join(part for part in (self.documentation, text) if part))
 
     def set_tools(self, tools):
         self._tool_state.registry = validate_tools(tools)
@@ -123,7 +122,7 @@ class Agent:
 
     def _update_system_message(self):
         content = (f"Today is {date.today().isoformat()}.\n\n<documentation>\n"
-                   + self.doc + "\n</documentation>\n\n" + self.instructions)
+                   + self.documentation + "\n</documentation>\n\n" + self.instructions)
         self.messages[:1] = [{"role": "system", "content": content}]
 
     def export_state(self):
@@ -168,6 +167,13 @@ class Agent:
                     entry["reasoning_details"] = message["reasoning_details"]
                 elif isinstance(message.get("reasoning"), str):
                     entry["reasoning"] = message["reasoning"]
+                if "provider_data" in message:
+                    data = message["provider_data"]
+                    if (not isinstance(data, dict) or data.get("provider") not in ("gemini", "anthropic")
+                            or not isinstance(data.get("content"), list)
+                            or not all(isinstance(block, dict) for block in data["content"])):
+                        raise ValueError("Invalid saved provider content.")
+                    entry["provider_data"] = data
             restored.append(entry)
         for result in results.values():
             if (not isinstance(result, dict) or not isinstance(result.get("sql"), str)
@@ -199,17 +205,14 @@ class Agent:
                 "These results are already saved and remain available even if older chat was trimmed. "
                 "Use query(result_id=...) or plot.sql to read them locally instead of paying to fetch "
                 "the same data again. Local SQL reads the table data using SQLite.\n" + json.dumps(saved)})
-        payload = {
-            "model": self.model, "messages": messages,
-            "max_tokens": 8192,
-            "reasoning": {"effort": os.environ.get("MULEBOT_REASONING_EFFORT", "medium"), "exclude": False},
-        }
-        if self._tool_state.registry:
-            payload.update(tools=self._tool_state.definitions, tool_choice=tool_choice, parallel_tool_calls=False)
+        url, headers, payload = build_request(
+            self.provider, self.model, self.model_api_key, messages, self._tool_state.definitions,
+            tool_choice, os.environ.get("MULEBOT_REASONING_EFFORT"),
+        )
         request = urllib.request.Request(
-            "https://openrouter.ai/api/v1/chat/completions",
+            url,
             data=json.dumps(payload).encode("utf-8"),
-            headers={"Authorization": f"Bearer {self.model_api_key}", "Content-Type": "application/json"},
+            headers=headers,
         )
         try:
             with urllib.request.urlopen(request, timeout=180) as response:
@@ -217,7 +220,7 @@ class Agent:
         except urllib.error.HTTPError as exc:
             detail = exc.read().decode("utf-8", errors="replace")[:2000]
             raise RuntimeError(self._redact(f"Model request failed ({exc.code}): {detail}")) from None
-        return result
+        return normalize_response(self.provider, result)
 
     def run(self, message):
         """Yield JSON-compatible events, independent of the CLI or browser UI."""
@@ -272,6 +275,8 @@ class Agent:
                     assistant["reasoning_details"] = response["reasoning_details"]
                 elif response.get("reasoning"):
                     assistant["reasoning"] = response["reasoning"]
+                if response.get("provider_data"):
+                    assistant["provider_data"] = response["provider_data"]
                 if not calls and not assistant["content"]:
                     raise RuntimeError("The model returned no answer. Try again or choose another model.")
                 self.messages.append(assistant)

{mulebot-0.2.0 → mulebot-0.3.0}/mulebot/cli.py

@@ -3,13 +3,15 @@ import json
 import sys
 
 from .agent import Agent
+from .providers import DEFAULT_MODELS
 
 
 def main():
     if hasattr(sys.stdout, "reconfigure"):
         sys.stdout.reconfigure(encoding="utf-8")
     parser = argparse.ArgumentParser(description="Chat with Datamule's APIs.")
-    parser.add_argument("--model", help="OpenRouter model; defaults to MULEBOT_MODEL.")
+    parser.add_argument("--provider", required=True, choices=DEFAULT_MODELS, help="Model provider.")
+    parser.add_argument("--model", help="Model ID; defaults to MULEBOT_MODEL or the provider's default.")
     parser.add_argument("--tables", help="Optional comma-separated table allowlist.")
     commands = parser.add_subparsers(dest="command", required=True)
     chat = commands.add_parser("chat", help="Print raw events as newline-delimited JSON.")
@@ -17,7 +19,7 @@ def main():
     serve = commands.add_parser("serve", help="Run the local browser interface.")
     serve.add_argument("--port", type=int, default=8000)
     args = parser.parse_args()
-    settings = {"model": args.model, "allowed_tables": None}
+    settings = {"provider": args.provider, "model": args.model, "allowed_tables": None}
     if args.tables is not None:
         settings["allowed_tables"] = [name.strip() for name in args.tables.split(",") if name.strip()]
     try:

mulebot-0.3.0/mulebot/providers.py

@@ -0,0 +1,171 @@
+"""Translate the shared agent history to each provider's HTTP API."""
+
+import json
+import os
+from urllib.parse import quote
+from uuid import uuid4
+
+
+DEFAULT_MODELS = {
+    "openrouter": "deepseek/deepseek-v4.1-flash-1017:nitro",
+    "openai": "gpt-5.4-mini",
+    "gemini": "gemini-3.8-flash",
+    "anthropic": "claude-sonnet-5-5",
+}
+KEY_VARIABLES = {
+    "openrouter": "OPENROUTER_API_KEY", "openai": "OPENAI_API_KEY",
+    "gemini": "GEMINI_API_KEY", "anthropic": "ANTHROPIC_API_KEY",
+}
+
+
+def validate_provider(provider):
+    if not isinstance(provider, str) or provider not in DEFAULT_MODELS:
+        raise ValueError("Choose a provider: " + ", ".join(DEFAULT_MODELS) + ".")
+    return provider
+
+
+def resolve_model(provider, model=None):
+    validate_provider(provider)
+    model = model if model is not None else os.environ.get("MULEBOT_MODEL") or DEFAULT_MODELS[provider]
+    if not isinstance(model, str) or not model.strip() or len(model) > 200 or any(c.isspace() for c in model.strip()):
+        raise ValueError("Enter a model ID of at most 200 characters.")
+    return model.strip()
+
+
+def credentials(provider):
+    validate_provider(provider)
+    missing = [name for name in ("DATAMULE_API_KEY", KEY_VARIABLES[provider]) if not os.environ.get(name, "").strip()]
+    if missing:
+        raise ValueError("Set " + " and ".join(missing) + ".")
+    return os.environ[KEY_VARIABLES[provider]]
+
+
+def _native_content(message, provider):
+    data = message.get("provider_data") or {}
+    return data.get("content") if data.get("provider") == provider else None
+
+
+def build_request(provider, model, key, messages, tools, tool_choice, reasoning_effort):
+    headers = {"Content-Type": "application/json"}
+    if provider in ("openrouter", "openai"):
+        history = []
+        for message in messages:
+            entry = {k: v for k, v in message.items() if k in ("role", "content", "tool_calls", "tool_call_id")}
+            if provider == "openrouter":
+                entry.update({k: v for k, v in message.items() if k in ("reasoning", "reasoning_details")})
+            history.append(entry)
+        payload = {"model": model, "messages": history}
+        if tools:
+            payload.update(tools=tools, tool_choice=tool_choice, parallel_tool_calls=False)
+        if provider == "openrouter":
+            payload.update(max_tokens=8192, reasoning={"effort": reasoning_effort or "medium", "exclude": False})
+            url = "https://openrouter.ai/api/v1/chat/completions"
+        else:
+            payload["max_completion_tokens"] = 8192
+            if reasoning_effort:
+                payload["reasoning_effort"] = reasoning_effort
+            url = "https://api.openai.com/v1/chat/completions"
+        headers["Authorization"] = f"Bearer {key}"
+        return url, headers, payload
+
+    system = "\n\n".join(message["content"] for message in messages if message["role"] == "system")
+    history, calls = [], {}
+    for message in messages:
+        role = message["role"]
+        if role == "system":
+            continue
+        for call in message.get("tool_calls", []):
+            calls[call["id"]] = call
+        native = _native_content(message, provider) if role == "assistant" else None
+        if provider == "anthropic":
+            target_role = "assistant" if role == "assistant" else "user"
+            if role == "tool":
+                content = [{"type": "tool_result", "tool_use_id": message["tool_call_id"], "content": message["content"]}]
+            elif native is not None:
+                content = native
+            else:
+                content = [{"type": "text", "text": message["content"]}] if message["content"] else []
+                content += [{"type": "tool_use", "id": call["id"], "name": call["function"]["name"],
+                             "input": json.loads(call["function"]["arguments"])} for call in message.get("tool_calls", [])]
+        else:
+            target_role = "model" if role == "assistant" else "user"
+            if role == "tool":
+                call = calls[message["tool_call_id"]]
+                response = json.loads(message["content"])
+                result = {"name": call["function"]["name"], "response": response if isinstance(response, dict) else {"result": response}}
+                # Older Gemini versions don't assign function-call IDs.
+                if not call["id"].startswith("mulebot_"):
+                    result["id"] = call["id"]
+                content = [{"functionResponse": result}]
+            elif native is not None:
+                content = native
+            else:
+                content = [{"text": message["content"]}] if message["content"] else []
+                content += [{"functionCall": {"name": call["function"]["name"], "args": json.loads(call["function"]["arguments"])}}
+                            for call in message.get("tool_calls", [])]
+        field = "content" if provider == "anthropic" else "parts"
+        if content:
+            # Group parallel results into one user message, as required by Claude.
+            if history and history[-1]["role"] == target_role:
+                history[-1][field].extend(content)
+            else:
+                history.append({"role": target_role, field: list(content)})
+
+    if provider == "anthropic":
+        headers.update({"x-api-key": key, "anthropic-version": "2023-06-01"})
+        payload = {"model": model, "system": system, "messages": history, "max_tokens": 8192}
+        if tools:
+            payload["tools"] = [{"name": tool["function"]["name"], "description": tool["function"]["description"],
+                                 "input_schema": tool["function"]["parameters"]} for tool in tools]
+            payload["tool_choice"] = {"type": tool_choice}
+            if tool_choice == "auto":
+                payload["tool_choice"]["disable_parallel_tool_use"] = True
+        return "https://api.anthropic.com/v1/messages", headers, payload
+
+    headers["x-goog-api-key"] = key
+    payload = {"systemInstruction": {"parts": [{"text": system}]}, "contents": history,
+               "generationConfig": {"maxOutputTokens": 8192}}
+    if tools:
+        payload["tools"] = [{"functionDeclarations": [{"name": tool["function"]["name"],
+                              "description": tool["function"]["description"],
+                              "parametersJsonSchema": tool["function"]["parameters"]} for tool in tools]}]
+        payload["toolConfig"] = {"functionCallingConfig": {"mode": tool_choice.upper()}}
+    return f"https://generativelanguage.googleapis.com/v1beta/models/{quote(model, safe='')}:generateContent", headers, payload
+
+
+def normalize_response(provider, response):
+    if provider in ("openrouter", "openai") or response.get("error"):
+        return response
+    calls, text, reasoning = [], [], []
+    if provider == "anthropic":
+        content = response.get("content") or []
+        for block in content:
+            if block["type"] == "text":
+                text.append(block["text"])
+            elif block["type"] == "thinking":
+                reasoning.append(block["thinking"])
+            elif block["type"] == "tool_use":
+                calls.append({"id": block["id"], "type": "function", "function": {
+                    "name": block["name"], "arguments": json.dumps(block["input"])}})
+        usage = dict(response.get("usage") or {})
+        usage.update(prompt_tokens=usage.get("input_tokens", 0) + usage.get("cache_creation_input_tokens", 0) + usage.get("cache_read_input_tokens", 0),
+                     completion_tokens=usage.get("output_tokens", 0))
+        model, response_id = response.get("model"), response.get("id")
+    else:
+        candidate = (response.get("candidates") or [{}])[0]
+        content = (candidate.get("content") or {}).get("parts") or []
+        for part in content:
+            if "text" in part:
+                (reasoning if part.get("thought") else text).append(part["text"])
+            if "functionCall" in part:
+                call = part["functionCall"]
+                calls.append({"id": call.get("id") or "mulebot_" + uuid4().hex, "type": "function", "function": {
+                    "name": call["name"], "arguments": json.dumps(call.get("args") or {})}})
+        usage = dict(response.get("usageMetadata") or {})
+        usage.update(prompt_tokens=usage.get("promptTokenCount", 0),
+                     completion_tokens=usage.get("candidatesTokenCount", 0) + usage.get("thoughtsTokenCount", 0))
+        model, response_id = response.get("modelVersion"), response.get("responseId")
+    usage["total_tokens"] = usage["prompt_tokens"] + usage["completion_tokens"]
+    message = {"content": "\n\n".join(text), "tool_calls": calls, "reasoning": "\n\n".join(reasoning),
+               "provider_data": {"provider": provider, "content": content}}
+    return {"id": response_id, "model": model, "usage": usage, "choices": [{"message": message}]}

{mulebot-0.2.0 → mulebot-0.3.0}/mulebot/server.py

@@ -1,10 +1,10 @@
 import json
-import os
 import sys
 from contextlib import closing
 from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
 
-from .agent import Agent, DEFAULT_MODEL
+from .agent import Agent
+from .providers import credentials, resolve_model
 from .assets import resource
 
 
@@ -43,7 +43,7 @@ class Handler(BaseHTTPRequestHandler):
             "/static/style.css": ("static/style.css", "text/css; charset=utf-8"),
         }
         if self.path == "/api/config":
-            self.send_json(200, {"model": self.server.default_model})
+            self.send_json(200, {"provider": self.server.provider, "model": self.server.default_model})
         elif self.path in assets:
             name, content_type = assets[self.path]
             body = resource("ui", *name.split("/")).read_text(encoding="utf-8")
@@ -75,7 +75,7 @@ class Handler(BaseHTTPRequestHandler):
                 raise ValueError("Enter a message of at most 16,000 characters.")
             model = body.get("model", self.server.default_model)
             if not isinstance(model, str) or not model.strip() or len(model) > 200 or any(c.isspace() for c in model.strip()):
-                raise ValueError("Enter an OpenRouter model ID of at most 200 characters.")
+                raise ValueError("Enter a model ID of at most 200 characters.")
             bot = self.server.agent_factory(model=model.strip())
             bot.load_state(body.get("state"))
         except (ValueError, UnicodeError) as exc:
@@ -94,15 +94,18 @@ class Handler(BaseHTTPRequestHandler):
             pass
 
 
-def create_server(port=8000, model=None, allowed_tables=None):
+def create_server(provider, *, port=8000, model=None, allowed_tables=None):
+    default_model = resolve_model(provider, model)
+    credentials(provider)
     server = ThreadingHTTPServer(("127.0.0.1", port), Handler)
-    server.default_model = model or os.environ.get("MULEBOT_MODEL", DEFAULT_MODEL)
-    server.agent_factory = lambda **settings: Agent(allowed_tables=allowed_tables, **settings)
+    server.provider = provider
+    server.default_model = default_model
+    server.agent_factory = lambda **settings: Agent(provider, allowed_tables=allowed_tables, **settings)
     return server
 
 
-def serve(port=8000, model=None, allowed_tables=None):
-    with create_server(port, model, allowed_tables) as server:
+def serve(provider, *, port=8000, model=None, allowed_tables=None):
+    with create_server(provider, port=port, model=model, allowed_tables=allowed_tables) as server:
         print(f"Mulebot: http://localhost:{server.server_port}", file=sys.stderr, flush=True)
         try:
             server.serve_forever()

{mulebot-0.2.0 → mulebot-0.3.0}/mulebot/tools.py

@@ -1,5 +1,6 @@
 import json
 import math
+import os
 import re
 import sqlite3
 import urllib.error
@@ -115,8 +116,10 @@ TOOL_DEFINITIONS = [
 
 
 class Tools:
-    def __init__(self, api_key, allowed_tables=None):
-        self.api_key = api_key
+    def __init__(self, *, allowed_tables=None):
+        self.api_key = os.environ.get("DATAMULE_API_KEY", "")
+        if not self.api_key.strip():
+            raise ValueError("Set DATAMULE_API_KEY.")
         self.allowed_tables = None if allowed_tables is None else {t.lower() for t in allowed_tables}
         self.catalog = None
         self.results = {}

{mulebot-0.2.0 → mulebot-0.3.0}/mulebot/ui/static/app.js

@@ -25,6 +25,7 @@ const toolCalls = new Map();
 const runningTools = new Set();
 let busy = false;
 let selectedModel = '';
+let selectedProvider = '';
 let agentState = null;
 const CHAT_STORAGE = 'mulebot.chat';
 let currentTools = null;
@@ -56,7 +57,7 @@ function debugHistory() {
 
 function saveChat() {
   try {
-    sessionStorage.setItem(CHAT_STORAGE, JSON.stringify({model: selectedModel, state: agentState, events: chatHistory}));
+    sessionStorage.setItem(CHAT_STORAGE, JSON.stringify({provider: selectedProvider, model: selectedModel, state: agentState, events: chatHistory}));
   } catch {
     settingsStatus.textContent = 'Chat is available in this tab but could not be saved for refresh.';
   }
@@ -522,7 +523,7 @@ function saveModel(choice) {
   if (!choice || choice === CUSTOM_MODEL) return;
   if (choice === selectedModel) { showSelectedModel(); setBusy(false); return; }
   if (choice.length > 200 || /\s/.test(choice)) {
-    settingsStatus.textContent = 'Enter an OpenRouter model ID of at most 200 characters.';
+    settingsStatus.textContent = 'Enter a model ID of at most 200 characters.';
     return;
   }
   selectedModel = choice; showSelectedModel(); clearChat(); settingsStatus.textContent = 'New chat started';
@@ -586,10 +587,13 @@ setBusy(true);
 fetch('/api/config').then(async response => {
   if (!response.ok) throw new Error('Could not load chat.');
   const config = await response.json(); selectedModel = config.model;
+  selectedProvider = config.provider;
+  if (selectedProvider !== 'openrouter') model.querySelector('optgroup')?.remove();
+  settingsStatus.textContent = `Provider: ${selectedProvider}`;
   let saved = null;
   try { saved = JSON.parse(sessionStorage.getItem(CHAT_STORAGE) || 'null'); }
   catch { settingsStatus.textContent = 'Could not restore the saved chat.'; }
-  if (saved && typeof saved.model === 'string' && Array.isArray(saved.events)) {
+  if (saved && saved.provider === selectedProvider && typeof saved.model === 'string' && Array.isArray(saved.events)) {
     selectedModel = saved.model; agentState = saved.state;
     for (const event of saved.events) apply(event);
     finishTools();

{mulebot-0.2.0 → mulebot-0.3.0}/mulebot/ui/templates/index.html

@@ -22,7 +22,7 @@
       </select>
     </label>
     <form id="custom-model-form" class="custom-model" hidden>
-      <label for="custom-model">OpenRouter model ID <input id="custom-model" placeholder="provider/model" maxlength="200" required></label>
+      <label for="custom-model">Model ID <input id="custom-model" placeholder="model-id" maxlength="200" required></label>
       <button id="use-custom-model" type="submit">Use model</button>
       <button id="cancel-custom-model" type="button">Cancel</button>
     </form>

{mulebot-0.2.0 → mulebot-0.3.0}/mulebot.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: mulebot
-Version: 0.2.0
+Version: 0.3.0
 Summary: A simple chatbot for Datamule's APIs
 Home-page: https://github.com/john-friedman/mulebot
 Author: John Friedman

{mulebot-0.2.0 → mulebot-0.3.0}/mulebot.egg-info/SOURCES.txt

@@ -7,6 +7,7 @@ mulebot/__main__.py
 mulebot/agent.py
 mulebot/assets.py
 mulebot/cli.py
+mulebot/providers.py
 mulebot/server.py
 mulebot/tool.py
 mulebot/tools.py

{mulebot-0.2.0 → mulebot-0.3.0}/setup.py

@@ -31,7 +31,7 @@ class Sdist(sdist):
 setup(
     name="mulebot",
     author="John Friedman",
-    version="0.2.0",
+    version="0.3.0",
     description="A simple chatbot for Datamule's APIs",
     url="https://github.com/john-friedman/mulebot",
     license="MIT",

mulebot-0.2.0/mulebot/__init__.py

@@ -1,4 +0,0 @@
-from .agent import Agent
-from .tool import Tool, tool
-
-__all__ = ["Agent", "Tool", "tool"]