"""Small, dependency-free client shared by Claude skills and the MCP bridge.""" from __future__ import annotations import argparse from dataclasses import dataclass import ipaddress import json import os from pathlib import Path import stat import urllib.error import urllib.parse import urllib.request MAX_RESPONSE_BYTES = 2 * 1024 * 1024 ENGINES = ('rag', 'graphrag') STATUSES = {'ok', 'no_answer', 'ambiguous', 'error'} class ClientError(Exception): """The message is a public reason code, never an exception from the transport.""" class NoRedirect(urllib.request.HTTPRedirectHandler): def redirect_request(self, req, fp, code, msg, headers, newurl): return None @dataclass(frozen=True) class Config: url: str key: str timeout: float = 300 allow_http: bool = False @classmethod def load(cls): values = {} filename = os.getenv('ZP_AGENT_CONFIG', '').strip() if filename: try: path = Path(filename).expanduser() if os.name == 'posix' and stat.S_IMODE(path.stat().st_mode) & 0o077: raise ClientError('config_permissions_require_0600') values = json.loads(path.read_text()) if not isinstance(values, dict): raise ValueError() except ClientError: raise except (OSError, ValueError): raise ClientError('invalid_configuration') from None try: config = cls(os.getenv('ZP_AGENT_URL', values.get('backend_url', '')).strip().rstrip('/'), os.getenv('ZP_AGENT_API_KEY', values.get('api_key', '')).strip(), float(os.getenv('ZP_AGENT_TIMEOUT', values.get('timeout_seconds', 300))), os.getenv('ZP_AGENT_ALLOW_HTTP', '').lower() in {'1', 'true'}) config.validate() return config except (ValueError, TypeError, AttributeError): raise ClientError('invalid_configuration') from None def validate(self): try: url = urllib.parse.urlsplit(self.url) if (url.scheme not in {'http', 'https'} or not url.hostname or url.username is not None or url.password is not None or url.query or url.fragment or any(c.isspace() for c in self.url)): raise ValueError() url.port if url.scheme == 'http': local = url.hostname == 'localhost' try: local = local or ipaddress.ip_address(url.hostname).is_loopback except ValueError: pass if not local and not self.allow_http: raise ValueError() if not self.key or '\r' in self.key or '\n' in self.key or not 1 <= self.timeout <= 600: raise ValueError() except (ValueError, TypeError): raise ClientError('invalid_configuration') from None def error_result(reason): messages={ 'unauthorized':'Prístup nebol autorizovaný. Skontroluj prihlásenie konektora a jeho konfiguráciu.', 'forbidden':'Pre tento nástroj nemáš potrebné oprávnenie.', 'invalid_query':'Otázku sa nepodarilo spracovať. Zadaj konkrétnu otázku a primeraný počet výsledkov.', 'model_unavailable':'Jazykový model momentálne nie je dostupný. Skús otázku neskôr.', 'answer_rejected':'Odpoveď sa nepodarilo spoľahlivo overiť podľa podkladov. ' 'Spresni osobu, typ práce alebo rok; neoverený text sa nezobrazí.', 'answer_capacity_exceeded':'Agent práve spracúva iné otázky. Skús to o chvíľu.', 'backend_busy':'Agent práve spracúva iné otázky. Skús to o chvíľu.', 'backend_timeout':'Vyhľadávanie neodpovedalo včas. Skús otázku zopakovať neskôr.', 'backend_unavailable':'Spojenie s backendom momentálne nie je dostupné. Skontroluj, či bežia služby a tunel.', 'retrieval_unavailable':'Vyhľadávací index momentálne nie je dostupný. Skús otázku neskôr.', } return {'status': 'error', 'reason': reason, 'accepted_for_delivery': False, 'answer': messages.get(reason,'Overenú odpoveď sa momentálne nepodarilo získať.'), 'sources': [], 'index_version': None, 'complete': False, 'truncated': False, 'truncation_reasons': []} def request_json(config: Config, endpoint: str, body=None): config.validate() request = urllib.request.Request(config.url + endpoint, data=json.dumps(body, ensure_ascii=False).encode() if body is not None else None, headers={'Authorization': 'Bearer ' + config.key, 'Accept': 'application/json', **({'Content-Type': 'application/json'} if body is not None else {})}) try: with urllib.request.build_opener(NoRedirect()).open(request, timeout=config.timeout) as response: raw = response.read(MAX_RESPONSE_BYTES + 1) except urllib.error.HTTPError as exc: code = exc.code exc.close() reason = {401: 'unauthorized', 403: 'forbidden', 404: 'backend_endpoint_not_found', 400: 'invalid_query', 422: 'invalid_request', 429: 'backend_busy', 503: 'backend_unavailable'}.get(code, 'backend_http_error') raise ClientError('redirect_not_allowed' if 300 <= code < 400 else reason) from None except (urllib.error.URLError, TimeoutError, OSError) as exc: timeout = isinstance(exc, TimeoutError) or isinstance(getattr(exc, 'reason', None), TimeoutError) raise ClientError('backend_timeout' if timeout else 'backend_unavailable') from None if len(raw) > MAX_RESPONSE_BYTES: raise ClientError('response_too_large') try: value = json.loads(raw) if not isinstance(value, dict): raise ValueError() return value except (ValueError, UnicodeError): raise ClientError('invalid_backend_response') from None def validate_metadata(value): if (not {'complete', 'truncated', 'truncation_reasons', 'index_version'} <= value.keys() or value.get('complete') is not None and type(value['complete']) is not bool or type(value.get('truncated')) is not bool or value['truncated'] and value.get('complete') is not False or not isinstance(value.get('truncation_reasons'), list) or not all(isinstance(x, str) for x in value['truncation_reasons']) or value['truncation_reasons'] and not value['truncated'] or value.get('index_version') is not None and not isinstance(value['index_version'], dict)): raise ClientError('invalid_backend_response') def query_backend(query: str, engine: str, *, mode='answer', limit=5, config=None): """Return only approved answer fields, or explicitly marked retrieval context.""" try: if (engine not in ENGINES or mode not in {'answer', 'retrieve'} or type(limit) is not int or not 1 <= limit <= 20 or not isinstance(query, str) or not 1 <= len(query.strip()) <= 500): raise ClientError('invalid_query') config = config or Config.load() endpoint = '/answer' if mode == 'answer' else '/' + engine body = {'query': query.strip(), 'limit': limit} if mode == 'answer' or engine == 'rag': body['include_diagnostics'] = False if mode == 'answer': body['engine'] = engine else: body['published_only'] = True value = request_json(config, endpoint, body) if not isinstance(value.get('status'), str) or value['status'] not in STATUSES: raise ClientError('invalid_backend_response') if value['status'] == 'error': # Raw model output, diagnostics and rejected text must not escape. reason = value.get('reason') safe_reasons = {'model_unavailable', 'answer_capacity_exceeded', 'answer_rejected', 'answer_unavailable', 'retrieval_unavailable'} raise ClientError(reason if isinstance(reason, str) and reason in safe_reasons else 'backend_answer_error' if mode == 'answer' else 'retrieval_unavailable') validate_metadata(value) fields = {'status', 'reason', 'index_version', 'complete', 'truncated', 'truncation_reasons'} if mode == 'answer': if value.get('accepted_for_delivery') is not True: raise ClientError('answer_rejected') if (not isinstance(value.get('answer'), str) or not value['answer'].strip() or not isinstance(value.get('source_urls'), list) or not all(isinstance(x, str) for x in value['source_urls']) or not isinstance(value.get('claims'), list)): raise ClientError('invalid_backend_response') if 'answer_coverage' in value: coverage = value['answer_coverage'] if (not isinstance(coverage, dict) or coverage.get('scope') not in ('fact','overview','assignment','list','count') or coverage.get('status') not in ('supported','partial','unknown','not_found','unavailable') or any(not isinstance(coverage.get(key), list) or any(not isinstance(item, str) for item in coverage[key]) for key in ['warnings','limitations'])): raise ClientError('invalid_backend_response') fields |= {'answer', 'source_urls', 'claims', 'count', 'accepted_for_delivery', 'answer_coverage'} else: if not isinstance(value.get('sources'), list): raise ClientError('invalid_backend_response') if engine == 'rag' and not isinstance(value.get('context'), str): raise ClientError('invalid_backend_response') if engine == 'graphrag' and (not isinstance(value.get('facts'), list) or not isinstance(value.get('paths'), list)): raise ClientError('invalid_backend_response') fields |= {'sources', 'context', 'facts', 'paths', 'aggregation', 'constraints', 'answer_candidates', 'instructions'} result = {key: item for key, item in value.items() if key in fields} return {**result, 'engine': engine, 'output_kind': 'verified_answer' if mode == 'answer' else 'retrieval_context', **({'requires_final_answer_verification': True} if mode == 'retrieve' else {})} except ClientError as exc: return error_result(str(exc)) def backend_health(config=None): try: value = request_json(config or Config.load(), '/health') tools = value.get('tools') if (type(value.get('ready')) is not bool or not isinstance(tools, dict) or any(not isinstance(tools.get(engine), dict) or type(tools[engine].get('ready')) is not bool for engine in ENGINES)): raise ClientError('invalid_backend_response') return {'status': 'ok' if value['ready'] else 'error', 'ready': value['ready'], 'index_version': value.get('index_version'), 'tools': { engine: {k: v for k, v in tools[engine].items() if k in {'ready', 'reason', 'index_version'}} for engine in ENGINES}} except ClientError as exc: return {**error_result(str(exc)), 'ready': False} def main(engine=None): parser = argparse.ArgumentParser(description=__doc__) parser.add_argument('query', nargs='?') if engine is None: parser.add_argument('--engine', choices=ENGINES, default='rag') parser.add_argument('--mode', choices=['answer', 'retrieve'], default='answer') parser.add_argument('--limit', type=int, default=5) parser.add_argument('--health', action='store_true') args = parser.parse_args() result = backend_health() if args.health else query_backend(args.query, engine or args.engine, mode=args.mode, limit=args.limit) print(json.dumps(result, ensure_ascii=False, separators=(',', ':'))) return 1 if result.get('status') == 'error' else 0 if __name__ == '__main__': raise SystemExit(main())