ServerKit/backend/app/services/ai_runtime.py

"""Adapter around Prompture drivers: isolated fallback and per-call accounting."""
import time

from app.services import ai_connections, ai_management


def managed_driver(driver, connection, model, recorder, max_cost, row):
    from prompture.drivers.base import Driver
    from prompture.exceptions import BudgetExceededError

    class ManagedDriver(Driver):
        def __init__(self):
            self.delegate = driver
            self.connection = connection
            self.model = model
            self.conversation = None
            self.tried = {(connection.id, model)}
            self._sync_flags()

        def _sync_flags(self):
            for name in ('supports_messages', 'supports_tool_use', 'supports_streaming',
                         'supports_streaming_tool_use', 'supports_vision', 'supports_json_mode', 'supports_json_schema'):
                setattr(self, name, getattr(self.delegate, name, False))

        def _switch(self, candidate, reason):
            target = ai_connections.get_connection(candidate['connection_id'])
            replacement = ai_connections.build_driver(target, candidate['model'])
            self.delegate, self.connection, self.model = replacement, target, candidate['model']
            self.tried.add((target.id, self.model))
            self._sync_flags()
            recorder.connection_id = target.id
            recorder.model = ai_connections.model_name(target, self.model)
            recorder.run.reason = reason
            recorder.switched = True
            row.connection_id, row.model_name = target.id, recorder.model
            if self.conversation is not None:
                self.conversation._model_name = recorder.model

        def _fallback(self, error):
            if (not recorder.config['fallback_enabled'] or recorder.output_started
                    or isinstance(error, (BudgetExceededError, ValueError, PermissionError))):
                return False
            for candidate in recorder.config['fallbacks']:
                key = candidate['connection_id'], candidate['model']
                if key in self.tried:
                    continue
                self.tried.add(key)
                try:
                    self._switch(candidate, 'Configured fallback after provider failure; no output or tools replayed')
                    return True
                except Exception:
                    # Failed construction is visible without leaking credentials.
                    recorder.attempts.append({**candidate, 'status': 'unavailable', 'cost_source': 'unknown',
                                              'cost': 0, 'total_tokens': 0, 'duration_ms': 0})
            return False

        def degrade(self):
            if recorder.config['budget_policy'] != 'degrade' or not max_cost:
                return
            if recorder.initial.get('cost', 0) < max_cost * 0.8:
                return
            rates = ai_management.model_info(self.connection, self.model)['pricing']
            if not rates:
                return
            for candidate in recorder.config['fallbacks']:
                target = ai_connections.get_connection(candidate['connection_id'])
                cheaper = ai_management.model_info(target, candidate['model'])['pricing']
                if cheaper and sum(cheaper.values()) < sum(rates.values()):
                    self._switch(candidate, 'Lower-cost model selected before the turn at 80% of the conversation allowance')
                    return

        def _options(self, options):
            # Keep SDK JSON/schema options, while enforcing administrator controls.
            return {**options, **ai_management.generation_options(recorder.config, self.connection, self.model)}

        def _meta(self, raw):
            from prompture.infra.tracker import UsageTracker
            import math
            raw = raw if isinstance(raw, dict) else {}
            meta = {}
            for key in ('prompt_tokens', 'completion_tokens', 'total_tokens', 'cached_prompt_tokens', 'cache_creation_tokens', 'cost'):
                value = raw.get(key, 0)
                meta[key] = value if isinstance(value, (int, float)) and math.isfinite(value) and value >= 0 else 0
            if not meta['total_tokens']:
                meta['total_tokens'] = meta['prompt_tokens'] + meta['completion_tokens']
            pricing = ai_management.model_info(self.connection, self.model)['pricing']
            reported = raw.get('cost')
            if (raw.get('cost_source') == 'provider_reported' and isinstance(reported, (int, float))
                    and not isinstance(reported, bool) and math.isfinite(reported) and reported >= 0):
                meta['cost_source'] = 'provider_reported'
            elif pricing is not None:
                if not meta['cost']:
                    meta['cost'] = UsageTracker.calculate_cost(ai_connections.model_name(self.connection, self.model), meta['prompt_tokens'], meta['completion_tokens'])
                meta['cost_source'] = 'estimated'
            else:
                meta['cost_source'] = 'estimated' if meta['cost'] > 0 else 'unknown'
            return meta

        def _call(self, method, *args):
            while True:
                recorder.check(max_cost)
                started = time.monotonic()
                try:
                    response = getattr(self.delegate, method)(*args[:-1], self._options(args[-1]))
                except Exception as error:
                    recorder.record({}, (time.monotonic() - started) * 1000, error=True)
                    if self._fallback(error):
                        continue
                    raise
                response = {**response, 'meta': self._meta(response.get('meta', {}))}
                recorder.output_started = True
                recorder.record(response['meta'], (time.monotonic() - started) * 1000)
                return response

        def _stream(self, method, *args):
            while True:
                recorder.check(max_cost)
                started = time.monotonic()
                meta = {}
                try:
                    for event in getattr(self.delegate, method)(*args[:-1], self._options(args[-1])):
                        recorder.output_started = True  # conservative: never replay any delivered event
                        if isinstance(event, dict) and 'meta' in event:
                            meta = self._meta(event['meta'])
                            event = {**event, 'meta': meta}
                        elif getattr(event, 'event_type', '') == 'message_stop':
                            meta = self._meta(event.usage)
                            event.usage.update(meta)
                        yield event
                except Exception as error:
                    recorder.record(meta, (time.monotonic() - started) * 1000, error=True)
                    if self._fallback(error):
                        continue
                    raise
                except GeneratorExit:
                    recorder.record(meta, (time.monotonic() - started) * 1000, error=True)
                    raise
                else:
                    recorder.record(meta, (time.monotonic() - started) * 1000)
                return

        def generate(self, prompt, options):
            return self._call('generate', prompt, options)

        def generate_messages(self, messages, options):
            return self._call('generate_messages', messages, options)

        def generate_messages_with_tools(self, messages, tools, options):
            return self._call('generate_messages_with_tools', messages, tools, options)

        def generate_messages_stream(self, messages, options):
            yield from self._stream('generate_messages_stream', messages, options)

        def generate_messages_with_tools_stream(self, messages, tools, options):
            yield from self._stream('generate_messages_with_tools_stream', messages, tools, options)

    wrapped = ManagedDriver()
    wrapped.degrade()
    return wrapped


EXTRACTION_SCHEMA = {'type': 'object', 'properties': {
    'summary': {'type': 'string'}, 'facts': {'type': 'array', 'items': {'type': 'string'}},
    'warnings': {'type': 'array', 'items': {'type': 'string'}},
}, 'required': ['summary', 'facts', 'warnings'], 'additionalProperties': False}


def ask(conv, message, workflow='chat'):
    if workflow == 'extract':
        result = conv.ask_for_json(message, EXTRACTION_SCHEMA, ai_cleanup=False)
        return result['json_string']
    return conv.ask(message)


def ask_live(conv, message, workflow='chat'):
    if workflow == 'extract':
        from prompture.agents.live_events import TextDelta, TurnComplete
        yield TextDelta(text=ask(conv, message, workflow))
        yield TurnComplete(usage=conv.usage)
    else:
        yield from conv.ask_live(message)