main
  1#!/usr/bin/env -S uv run --script
  2# /// script
  3# dependencies = [
  4#     "requests",
  5#     "PyYAML",
  6# ]
  7# ///
  8#
  9# Generate aichat config.yaml from config.yaml.in template.
 10# - Resolves passage:: secrets
 11# - Fetches available models from running services (Ollama, Gemini, etc.)
 12# - Uses REST API for Gemini (no grpc/native extensions needed)
 13
 14import os
 15import os.path
 16import subprocess
 17import socket
 18import sys
 19from typing import Any, Dict, List, Optional
 20
 21import requests
 22import yaml
 23import urllib.parse as urlparse
 24
 25
 26def debug(msg: str):
 27    print(f"[DEBUG] {msg}", file=sys.stderr)
 28
 29
 30def check_running(api_base: str, timeout: float = 0.5) -> bool:
 31    """Quickly check if a service is accessible at the given host and port."""
 32    url = urlparse.urlparse(api_base)
 33    port = url.port or (80 if url.scheme == "http" else 443)
 34    debug(f"Checking if {url.hostname}:{port} is running (api_base={api_base})")
 35    try:
 36        sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
 37        sock.settimeout(timeout)
 38        result = sock.connect_ex((url.hostname, port))
 39        sock.close()
 40        debug(f"Socket connect_ex result for {url.hostname}:{port}: {result}")
 41        return result == 0
 42    except Exception as e:
 43        debug(f"Exception in check_running: {e}")
 44        return False
 45
 46
 47def load_config(config_path: str) -> Dict[str, Any]:
 48    debug(f"Loading config from {config_path}")
 49    with open(config_path, "r") as file:
 50        config = yaml.safe_load(file)
 51        debug(f"Loaded config: {config}")
 52        return config
 53
 54
 55def get_passageword(passage_path: str) -> Optional[str]:
 56    """Retrieve password from passage using the given path."""
 57    debug(f"get_passageword called for passage_path={passage_path}")
 58    try:
 59        result = subprocess.run(
 60            ["passage", "show", passage_path], capture_output=True, text=True, check=True
 61        )
 62        passageword = result.stdout.splitlines()[0]
 63        debug(f"Password retrieved from passage for {passage_path}")
 64        return passageword
 65    except FileNotFoundError:
 66        print(
 67            "Error: 'passage' command not found. Is passage installed and in your PATH?",
 68            file=sys.stderr,
 69        )
 70        debug("'passage' command not found")
 71        return None
 72    except subprocess.CalledProcessError as e:
 73        print(f"Error running passage show {passage_path}: {e.stderr}", file=sys.stderr)
 74        debug(f"subprocess.CalledProcessError in get_passageword: {e}")
 75        return None
 76    except IndexError:
 77        print(f"Error: 'passage show {passage_path}' returned empty output.", file=sys.stderr)
 78        debug(f"IndexError: passage show {passage_path} returned empty output")
 79        return None
 80
 81
 82def resolve_api_key(api_key_config: Optional[str]) -> Optional[str]:
 83    """Resolve API key, fetching from passage if needed."""
 84    if not api_key_config:
 85        return None
 86    if api_key_config.startswith("passage::"):
 87        passage_path = api_key_config.split("::", 1)[1]
 88        debug(f"Retrieving API key from passage at {passage_path}")
 89        return get_passageword(passage_path)
 90    return api_key_config
 91
 92
 93def get_models_openai(api_base: str, api_key: Optional[str]) -> List[Dict[str, str]]:
 94    """Query OpenAI-compatible /models endpoint."""
 95    debug(f"get_models_openai: api_base={api_base}")
 96
 97    headers = {}
 98    if api_key:
 99        headers["Authorization"] = f"Bearer {api_key}"
100
101    if not api_base.endswith("/"):
102        api_base = api_base + "/"
103
104    models_url = f"{api_base}models"
105    debug(f"Querying models endpoint: {models_url}")
106
107    try:
108        response = requests.get(models_url, headers=headers, timeout=10)
109        debug(f"HTTP GET {models_url} status_code={response.status_code}")
110        response.raise_for_status()
111        data = response.json()
112
113        models = data.get("data", [])
114        debug(f"Found {len(models)} models")
115        return [
116            {"name": model.get("id"), "description": model.get("id")}
117            for model in models
118        ]
119    except Exception as e:
120        print(f"Error querying {api_base}: {str(e)}", file=sys.stderr)
121        debug(f"Exception in get_models_openai: {e}")
122        return []
123
124
125def get_gemini_models(api_key: Optional[str]) -> List[Dict[str, str]]:
126    """Query Google's Gemini API via REST and return available models.
127
128    Uses the REST API directly instead of the google-generativeai SDK
129    to avoid grpc/native extension dependencies.
130    """
131    debug(f"get_gemini_models called with api_key={'***' if api_key else None}")
132    if not api_key:
133        print("Error: API key is required for Google Gemini API", file=sys.stderr)
134        return []
135
136    # Gemini REST API endpoint for listing models
137    url = "https://generativelanguage.googleapis.com/v1beta/models"
138    headers = {"x-goog-api-key": api_key}
139
140    try:
141        debug(f"Querying Gemini REST API: {url}")
142        response = requests.get(url, headers=headers, timeout=10)
143        debug(f"HTTP GET {url} status_code={response.status_code}")
144        response.raise_for_status()
145        data = response.json()
146
147        # Filter for gemini models and extract name/description
148        models = data.get("models", [])
149        gemini_models = [
150            {
151                "name": model.get("name", "").split("/")[-1],  # "models/gemini-pro" -> "gemini-pro"
152                "description": model.get("description", model.get("displayName", "")),
153            }
154            for model in models
155            if "gemini" in model.get("name", "").lower()
156        ]
157        debug(f"Found {len(gemini_models)} Gemini models")
158        return gemini_models
159
160    except requests.exceptions.HTTPError as e:
161        print(f"HTTP error querying Gemini API: {e}", file=sys.stderr)
162        debug(f"HTTPError in get_gemini_models: {e}")
163        return []
164    except Exception as e:
165        print(f"Error querying Google Gemini API: {str(e)}", file=sys.stderr)
166        debug(f"Exception in get_gemini_models: {e}")
167        return []
168
169
170def process_client(client: Dict[str, Any]) -> Optional[Dict[str, Any]]:
171    """Process a single client configuration."""
172    updated_client = client.copy()
173    debug(f"Processing client: {updated_client.get('name', 'unnamed')}")
174
175    # Resolve API key from passage if needed
176    api_key_config = updated_client.get("api_key")
177    actual_api_key = resolve_api_key(api_key_config)
178    updated_client["api_key"] = actual_api_key
179
180    client_type = updated_client.get("type")
181
182    # OpenAI-compatible clients (Ollama, Groq, etc.)
183    if client_type == "openai-compatible":
184        if not updated_client.get("models"):
185            api_base = updated_client.get("api_base")
186            if api_base:
187                if not check_running(api_base):
188                    debug(f"{api_base} not running, skipping client")
189                    return None
190                fetched_models = get_models_openai(api_base, actual_api_key)
191                updated_client["models"] = fetched_models if fetched_models else []
192            else:
193                updated_client["models"] = []
194
195    # Google Gemini clients
196    elif client_type == "gemini":
197        if not updated_client.get("models"):
198            fetched_models = get_gemini_models(actual_api_key)
199            updated_client["models"] = fetched_models if fetched_models else []
200
201    return updated_client
202
203
204def main():
205    # Support reading from script directory or ~/.config/aichat/
206    script_dir = os.path.dirname(os.path.abspath(__file__))
207    config_path = os.path.join(script_dir, "config.yaml.in")
208
209    if not os.path.exists(config_path):
210        config_path = os.path.expanduser("~/.config/aichat/config.yaml.in")
211
212    debug(f"main: config_path={config_path}")
213    config_data = load_config(config_path)
214
215    if "clients" in config_data:
216        updated_clients = []
217        debug(f"main: found {len(config_data.get('clients', []))} clients")
218
219        for client in config_data.get("clients", []):
220            processed = process_client(client)
221            if processed is not None:
222                updated_clients.append(processed)
223                debug(f"main: client {client.get('name', 'unnamed')} processed")
224
225        config_data["clients"] = updated_clients
226
227    # Output final configuration as YAML
228    debug("main: dumping config_data as YAML")
229    print(yaml.dump(config_data, default_flow_style=False, sort_keys=False))
230
231
232if __name__ == "__main__":
233    main()