fix(consult): preserve streaming across disconnects
This commit is contained in:
Jesse_Chen
@@ -0,0 +1,280 @@
|
||||
begin;
|
||||
|
||||
alter table public.consultation_requests
|
||||
add column if not exists session_id uuid references public.chat_sessions(id) on delete set null,
|
||||
add column if not exists response_message jsonb;
|
||||
|
||||
alter table public.consultation_requests
|
||||
drop constraint if exists consultation_requests_response_message_check;
|
||||
alter table public.consultation_requests
|
||||
add constraint consultation_requests_response_message_check
|
||||
check (response_message is null or jsonb_typeof(response_message) = 'object');
|
||||
|
||||
create index if not exists consultation_requests_user_session_created_idx
|
||||
on public.consultation_requests (user_id, session_id, created_at desc);
|
||||
|
||||
create index if not exists consultation_requests_user_status_created_idx
|
||||
on public.consultation_requests (user_id, status, created_at desc);
|
||||
|
||||
drop function if exists public.reserve_consultation_usage(uuid, text, text, integer);
|
||||
|
||||
create or replace function public.reserve_consultation_usage(
|
||||
p_user_id uuid,
|
||||
p_request_id text,
|
||||
p_session_id uuid,
|
||||
p_requested_model_id text,
|
||||
p_credit_cost integer
|
||||
)
|
||||
returns table(
|
||||
success boolean,
|
||||
reservation_id uuid,
|
||||
source text,
|
||||
credits integer,
|
||||
subscription_id uuid,
|
||||
reason text,
|
||||
retry_after_seconds integer
|
||||
)
|
||||
language plpgsql
|
||||
security definer
|
||||
set search_path = ''
|
||||
as $$
|
||||
declare
|
||||
v_authorization record;
|
||||
begin
|
||||
if p_user_id is null or p_session_id is null or btrim(coalesce(p_request_id, '')) = '' then
|
||||
return query select false, null::uuid, null::text, null::integer,
|
||||
null::uuid, 'invalid_request'::text, null::integer;
|
||||
return;
|
||||
end if;
|
||||
|
||||
perform pg_advisory_xact_lock(hashtextextended(p_user_id::text || ':' || btrim(p_request_id), 0));
|
||||
|
||||
if not exists (
|
||||
select 1
|
||||
from public.chat_sessions as session
|
||||
where session.id = p_session_id
|
||||
and session.user_id = p_user_id
|
||||
and session.session_type = 'consultation'
|
||||
) then
|
||||
return query select false, null::uuid, null::text,
|
||||
(select profile.credits from public.profiles as profile where profile.id = p_user_id),
|
||||
null::uuid, 'session_missing'::text, null::integer;
|
||||
return;
|
||||
end if;
|
||||
|
||||
if exists (
|
||||
select 1
|
||||
from public.consultation_requests as request
|
||||
where request.user_id = p_user_id
|
||||
and request.request_id = btrim(p_request_id)
|
||||
) then
|
||||
return query select false, null::uuid, null::text,
|
||||
(select profile.credits from public.profiles as profile where profile.id = p_user_id),
|
||||
null::uuid, 'request_conflict'::text, null::integer;
|
||||
return;
|
||||
end if;
|
||||
|
||||
select * into v_authorization
|
||||
from public.authorize_usage(
|
||||
p_user_id,
|
||||
'chat.standard',
|
||||
p_requested_model_id,
|
||||
btrim(p_request_id),
|
||||
p_credit_cost
|
||||
);
|
||||
|
||||
if coalesce(v_authorization.success, false) then
|
||||
insert into public.consultation_requests(user_id, request_id, session_id, status)
|
||||
values (p_user_id, btrim(p_request_id), p_session_id, 'reserved');
|
||||
end if;
|
||||
|
||||
return query select
|
||||
v_authorization.success,
|
||||
v_authorization.reservation_id,
|
||||
v_authorization.source,
|
||||
v_authorization.credits,
|
||||
v_authorization.subscription_id,
|
||||
v_authorization.reason,
|
||||
v_authorization.retry_after_seconds;
|
||||
end;
|
||||
$$;
|
||||
|
||||
revoke all on function public.reserve_consultation_usage(uuid, text, uuid, text, integer)
|
||||
from public, anon, authenticated;
|
||||
grant execute on function public.reserve_consultation_usage(uuid, text, uuid, text, integer)
|
||||
to service_role;
|
||||
do $$ begin
|
||||
if exists(select 1 from pg_roles where rolname = 'admin_runtime') then
|
||||
grant execute on function public.reserve_consultation_usage(uuid, text, uuid, text, integer)
|
||||
to admin_runtime;
|
||||
end if;
|
||||
end $$;
|
||||
|
||||
create or replace function public.cancel_consultation_credit(p_user_id uuid, p_request_id text)
|
||||
returns table(success boolean, credits integer, error_code text)
|
||||
language plpgsql
|
||||
security definer
|
||||
set search_path = ''
|
||||
as $$
|
||||
declare
|
||||
v_request public.consultation_requests%rowtype;
|
||||
v_release record;
|
||||
v_balance integer;
|
||||
begin
|
||||
if p_user_id is null or btrim(coalesce(p_request_id, '')) = '' then
|
||||
return query select false, null::integer, 'invalid_request'::text;
|
||||
return;
|
||||
end if;
|
||||
|
||||
perform pg_advisory_xact_lock(hashtextextended(p_user_id::text || ':' || btrim(p_request_id), 0));
|
||||
|
||||
select request.* into v_request
|
||||
from public.consultation_requests as request
|
||||
where request.user_id = p_user_id
|
||||
and request.request_id = btrim(p_request_id)
|
||||
for update;
|
||||
|
||||
if not found then
|
||||
return query select false, null::integer, 'request_missing'::text;
|
||||
return;
|
||||
end if;
|
||||
|
||||
select profile.credits into v_balance
|
||||
from public.profiles as profile
|
||||
where profile.id = p_user_id;
|
||||
|
||||
if v_request.status = 'completed' then
|
||||
return query select false, v_balance, 'request_completed'::text;
|
||||
return;
|
||||
end if;
|
||||
if v_request.status = 'cancelled' then
|
||||
return query select true, v_balance, null::text;
|
||||
return;
|
||||
end if;
|
||||
if v_request.status <> 'reserved' then
|
||||
return query select false, v_balance, 'invalid_request_status'::text;
|
||||
return;
|
||||
end if;
|
||||
|
||||
select * into v_release
|
||||
from public.release_usage(p_user_id, btrim(p_request_id), 'consultation_cancelled');
|
||||
if not coalesce(v_release.success, false) then
|
||||
return query select false, v_release.credits, v_release.error_code;
|
||||
return;
|
||||
end if;
|
||||
|
||||
update public.consultation_requests
|
||||
set status = 'cancelled', updated_at = clock_timestamp()
|
||||
where user_id = p_user_id
|
||||
and request_id = btrim(p_request_id)
|
||||
and status = 'reserved';
|
||||
|
||||
return query select true, v_release.credits, null::text;
|
||||
end;
|
||||
$$;
|
||||
|
||||
revoke all on function public.cancel_consultation_credit(uuid, text)
|
||||
from public, anon, authenticated;
|
||||
grant execute on function public.cancel_consultation_credit(uuid, text)
|
||||
to service_role;
|
||||
do $$ begin
|
||||
if exists(select 1 from pg_roles where rolname = 'admin_runtime') then
|
||||
grant execute on function public.cancel_consultation_credit(uuid, text)
|
||||
to admin_runtime;
|
||||
end if;
|
||||
end $$;
|
||||
|
||||
create or replace function public.complete_consultation_response(
|
||||
p_user_id uuid,
|
||||
p_request_id text,
|
||||
p_session_id uuid,
|
||||
p_response_message jsonb,
|
||||
p_actual_usage jsonb
|
||||
)
|
||||
returns table(success boolean, credits integer, error_code text)
|
||||
language plpgsql
|
||||
security definer
|
||||
set search_path = ''
|
||||
as $$
|
||||
declare
|
||||
v_request public.consultation_requests%rowtype;
|
||||
v_settlement record;
|
||||
v_balance integer;
|
||||
begin
|
||||
if p_user_id is null or p_session_id is null or btrim(coalesce(p_request_id, '')) = '' then
|
||||
return query select false, null::integer, 'invalid_request'::text;
|
||||
return;
|
||||
end if;
|
||||
if jsonb_typeof(p_response_message) <> 'object'
|
||||
or p_response_message->>'role' <> 'assistant'
|
||||
or btrim(coalesce(p_response_message->>'text', '')) = '' then
|
||||
return query select false, null::integer, 'invalid_response_message'::text;
|
||||
return;
|
||||
end if;
|
||||
|
||||
perform pg_advisory_xact_lock(hashtextextended(p_user_id::text || ':' || btrim(p_request_id), 0));
|
||||
|
||||
select request.* into v_request
|
||||
from public.consultation_requests as request
|
||||
where request.user_id = p_user_id
|
||||
and request.request_id = btrim(p_request_id)
|
||||
and request.session_id = p_session_id
|
||||
for update;
|
||||
|
||||
if not found then
|
||||
return query select false, null::integer, 'request_missing'::text;
|
||||
return;
|
||||
end if;
|
||||
if v_request.status = 'cancelled' then
|
||||
return query select false, null::integer, 'request_cancelled'::text;
|
||||
return;
|
||||
end if;
|
||||
if v_request.status = 'completed' then
|
||||
select profile.credits into v_balance from public.profiles as profile where profile.id = p_user_id;
|
||||
return query select v_request.response_message = p_response_message, v_balance,
|
||||
case when v_request.response_message = p_response_message then null::text else 'response_conflict'::text end;
|
||||
return;
|
||||
end if;
|
||||
|
||||
update public.chat_sessions as session
|
||||
set messages = session.messages || jsonb_build_array(p_response_message),
|
||||
updated_at = clock_timestamp()
|
||||
where session.id = p_session_id
|
||||
and session.user_id = p_user_id
|
||||
and session.session_type = 'consultation';
|
||||
if not found then
|
||||
return query select false, null::integer, 'session_missing'::text;
|
||||
return;
|
||||
end if;
|
||||
|
||||
update public.consultation_requests
|
||||
set response_message = p_response_message,
|
||||
updated_at = clock_timestamp()
|
||||
where user_id = p_user_id and request_id = btrim(p_request_id);
|
||||
|
||||
select * into v_settlement
|
||||
from public.complete_usage(p_user_id, btrim(p_request_id), coalesce(p_actual_usage, '{}'::jsonb));
|
||||
if not coalesce(v_settlement.success, false) then
|
||||
raise exception 'consultation_usage_settlement_failed:%', coalesce(v_settlement.error_code, 'unknown');
|
||||
end if;
|
||||
|
||||
update public.consultation_requests
|
||||
set status = 'completed', updated_at = clock_timestamp()
|
||||
where user_id = p_user_id and request_id = btrim(p_request_id);
|
||||
|
||||
return query select true, v_settlement.credits, null::text;
|
||||
end;
|
||||
$$;
|
||||
|
||||
revoke all on function public.complete_consultation_response(uuid, text, uuid, jsonb, jsonb)
|
||||
from public, anon, authenticated;
|
||||
grant execute on function public.complete_consultation_response(uuid, text, uuid, jsonb, jsonb)
|
||||
to service_role;
|
||||
do $$ begin
|
||||
if exists(select 1 from pg_roles where rolname = 'admin_runtime') then
|
||||
grant execute on function public.complete_consultation_response(uuid, text, uuid, jsonb, jsonb)
|
||||
to admin_runtime;
|
||||
end if;
|
||||
end $$;
|
||||
|
||||
commit;
|
||||
Reference in New Issue
Block a user