SELECT AI on Oracle 23ai and 26ai: Tracking Token Usage per Schema

-- Create the observability schema
CREATE USER ai_ops IDENTIFIED BY "&strong_password"
  DEFAULT TABLESPACE data
  QUOTA UNLIMITED ON data;

-- Required privileges
GRANT CREATE SESSION    TO ai_ops;
GRANT CREATE TABLE      TO ai_ops;
GRANT CREATE PROCEDURE  TO ai_ops;
GRANT CREATE SEQUENCE   TO ai_ops;
GRANT SELECT ON dba_cloud_ai_profile_attributes TO ai_ops;  -- model name in log
GRANT SELECT ON v$session                        TO ai_ops;  -- SID/SERIAL# capture
GRANT EXECUTE ON dbms_cloud_ai                   TO ai_ops;  -- LLM calls

-- Enable Resource Principal for AI_OPS
-- Eliminates the need for an explicit credential to access OCI GenAI
EXEC DBMS_CLOUD_ADMIN.ENABLE_RESOURCE_PRINCIPAL(username => 'AI_OPS');

-- Main usage tracking table
CREATE TABLE ai_ops.ai_token_log (
  log_id            NUMBER         GENERATED ALWAYS AS IDENTITY PRIMARY KEY,
  log_ts            TIMESTAMP      DEFAULT SYSTIMESTAMP NOT NULL,
  caller_schema     VARCHAR2(128)  NOT NULL,
  profile_name      VARCHAR2(128)  NOT NULL,
  action            VARCHAR2(20)   NOT NULL,
  prompt_text       CLOB,
  response_text     CLOB,
  tokens_input      NUMBER,
  tokens_output     NUMBER,
  tokens_total      NUMBER         GENERATED ALWAYS AS (tokens_input + tokens_output) VIRTUAL,
  model_name        VARCHAR2(256),
  duration_ms       NUMBER,
  session_sid       NUMBER,
  session_serial    NUMBER,
  client_identifier VARCHAR2(128),
  client_ip         VARCHAR2(128),
  error_code        NUMBER,
  error_message     VARCHAR2(4000),
  status            VARCHAR2(10)   DEFAULT 'OK' NOT NULL
    CONSTRAINT chk_status CHECK (status IN ('OK','ERROR'))
) TABLESPACE data;

-- Indexes to support consumption reports
CREATE INDEX ai_token_log_caller_ts_ix
  ON ai_ops.ai_token_log (caller_schema, TRUNC(log_ts));

CREATE INDEX ai_token_log_ts_ix
  ON ai_ops.ai_token_log (TRUNC(log_ts));

CREATE INDEX ai_token_log_action_ix
  ON ai_ops.ai_token_log (action, caller_schema);

-- Comments
COMMENT ON TABLE ai_ops.ai_token_log IS
  'Centralized log of all SELECT AI calls made through ai_ops.pkg_ai_gateway. '
  'Each row represents one call to DBMS_CLOUD_AI.GENERATE.';

COMMENT ON COLUMN ai_ops.ai_token_log.caller_schema IS
  'Schema that originated the call. Captured via SYS_CONTEXT — cannot be forged by the caller.';

COMMENT ON COLUMN ai_ops.ai_token_log.tokens_input IS
  'Input tokens (prompt + schema metadata) as reported by the OCI GenAI API.';

COMMENT ON COLUMN ai_ops.ai_token_log.tokens_output IS
  'Output tokens (LLM response) as reported by the OCI GenAI API.';

COMMENT ON COLUMN ai_ops.ai_token_log.tokens_total IS
  'Virtual column: sum of tokens_input + tokens_output.';

CREATE OR REPLACE PACKAGE ai_ops.pkg_ai_gateway AS

  -- Main function: LLM call with automatic logging
  FUNCTION generate (
    p_prompt       IN CLOB,
    p_profile_name IN VARCHAR2,
    p_action       IN VARCHAR2 DEFAULT 'runsql'
  ) RETURN CLOB;

  -- Helper procedure: log an error with no result
  PROCEDURE log_error (
    p_caller_schema IN VARCHAR2,
    p_profile_name  IN VARCHAR2,
    p_action        IN VARCHAR2,
    p_prompt        IN CLOB,
    p_error_code    IN NUMBER,
    p_error_message IN VARCHAR2
  );

END pkg_ai_gateway;
/

CREATE OR REPLACE PACKAGE BODY ai_ops.pkg_ai_gateway AS

  -- Token estimation by character count
  -- Approximation: 1 token ≈ 4 chars (English/Portuguese)
  -- For exact counts, replace with a call to the provider's tokenization API
  FUNCTION estimate_tokens (p_text IN CLOB) RETURN NUMBER IS
  BEGIN
    RETURN CEIL(DBMS_LOB.GETLENGTH(NVL(p_text, EMPTY_CLOB())) / 4);
  END estimate_tokens;

  FUNCTION generate (
    p_prompt       IN CLOB,
    p_profile_name IN VARCHAR2,
    p_action       IN VARCHAR2 DEFAULT 'runsql'
  ) RETURN CLOB IS

    l_result         CLOB;
    l_start_ts       TIMESTAMP := SYSTIMESTAMP;
    l_start_time     NUMBER    := DBMS_UTILITY.GET_TIME;
    l_duration_ms    NUMBER;
    l_tokens_in      NUMBER;
    l_tokens_out     NUMBER;
    l_caller_schema  VARCHAR2(128) := SYS_CONTEXT('USERENV', 'SESSION_USER');
    l_sid            NUMBER        := SYS_CONTEXT('USERENV', 'SID');
    l_serial         NUMBER;
    l_client_id      VARCHAR2(128) := SYS_CONTEXT('USERENV', 'CLIENT_IDENTIFIER');
    l_client_ip      VARCHAR2(128) := SYS_CONTEXT('USERENV', 'IP_ADDRESS');
    l_model_name     VARCHAR2(256);
    l_usage_json     JSON_OBJECT_T;
    l_json_response  CLOB;

  BEGIN
    -- Retrieve SERIAL# for the current session
    SELECT serial#
      INTO l_serial
      FROM v$session
     WHERE sid = l_sid
       AND rownum = 1;

    -- For 'chat', the response is the full API JSON (includes usage)
    -- For other actions, the response is the final result (SQL, text, etc.)
    IF p_action = 'chat' THEN
      l_json_response := DBMS_CLOUD_AI.GENERATE(
        prompt       => p_prompt,
        profile_name => p_profile_name,
        action       => p_action
      );

      -- Try to extract tokens from the usage field
      BEGIN
        l_usage_json  := JSON_OBJECT_T.PARSE(l_json_response);
        l_tokens_in   := l_usage_json.get_Object('usage').get_Number('prompt_tokens');
        l_tokens_out  := l_usage_json.get_Object('usage').get_Number('completion_tokens');
        l_result      := l_usage_json.get_String('content');
        IF l_result IS NULL THEN
          l_result := l_json_response; -- fallback: return the full JSON
        END IF;
      EXCEPTION
        WHEN OTHERS THEN
          -- JSON not in expected format; return everything and estimate tokens
          l_result     := l_json_response;
          l_tokens_in  := estimate_tokens(p_prompt);
          l_tokens_out := estimate_tokens(l_result);
      END;

    ELSE
      -- runsql, narrate, showsql: return result directly
      l_result     := DBMS_CLOUD_AI.GENERATE(
        prompt       => p_prompt,
        profile_name => p_profile_name,
        action       => p_action
      );
      -- Token estimation (replace with tokenization API call if needed)
      l_tokens_in  := estimate_tokens(p_prompt);
      l_tokens_out := estimate_tokens(l_result);
    END IF;

    -- Retrieve the model name configured in the profile
    BEGIN
      SELECT attribute_value
        INTO l_model_name
        FROM dba_cloud_ai_profile_attributes
       WHERE profile_name   = p_profile_name
         AND attribute_name = 'model'
         AND rownum = 1;
    EXCEPTION
      WHEN NO_DATA_FOUND THEN
        l_model_name := NULL;
    END;

    -- Calculate duration in milliseconds
    -- DBMS_UTILITY.GET_TIME returns hundredths of a second
    l_duration_ms := (DBMS_UTILITY.GET_TIME - l_start_time) * 10;

    -- Write log
    INSERT INTO ai_ops.ai_token_log (
      log_ts, caller_schema, profile_name, action,
      prompt_text, response_text,
      tokens_input, tokens_output,
      model_name, duration_ms,
      session_sid, session_serial,
      client_identifier, client_ip,
      status
    ) VALUES (
      l_start_ts, l_caller_schema, p_profile_name, p_action,
      p_prompt, l_result,
      l_tokens_in, l_tokens_out,
      l_model_name,
      l_duration_ms,
      l_sid, l_serial,
      l_client_id, l_client_ip,
      'OK'
    );

    COMMIT;
    RETURN l_result;

  EXCEPTION
    WHEN OTHERS THEN
      l_duration_ms := (DBMS_UTILITY.GET_TIME - l_start_time) * 10;
      log_error(
        p_caller_schema => l_caller_schema,
        p_profile_name  => p_profile_name,
        p_action        => p_action,
        p_prompt        => p_prompt,
        p_error_code    => SQLCODE,
        p_error_message => SQLERRM
      );
      RAISE;
  END generate;

  PROCEDURE log_error (
    p_caller_schema IN VARCHAR2,
    p_profile_name  IN VARCHAR2,
    p_action        IN VARCHAR2,
    p_prompt        IN CLOB,
    p_error_code    IN NUMBER,
    p_error_message IN VARCHAR2
  ) IS
    PRAGMA AUTONOMOUS_TRANSACTION;
  BEGIN
    INSERT INTO ai_ops.ai_token_log (
      log_ts, caller_schema, profile_name, action,
      prompt_text, tokens_input, tokens_output,
      error_code, error_message, status
    ) VALUES (
      SYSTIMESTAMP, p_caller_schema, p_profile_name, p_action,
      p_prompt, 0, 0,
      p_error_code, p_error_message, 'ERROR'
    );
    COMMIT;
  END log_error;

END pkg_ai_gateway;
/

-- Profile for APP_SALES
-- Context: clients, contracts and environments
BEGIN
  DBMS_CLOUD_AI.CREATE_PROFILE(
    profile_name => 'AI_PROFILE_SALES',
    attributes   => '{
      "provider"          : "oci",
      "credential_name"   : "OCI$RESOURCE_PRINCIPAL",
      "model"             : "cohere.command-r-plus-08-2024",
      "oci_compartment_id": "ocid1.compartment.oc1..xxxxxxxx",
      "temperature"       : 0,
      "comments"          : true,
      "object_list"       : [
        {"owner": "RREZENDE", "name": "EMPRESA"},
        {"owner": "RREZENDE", "name": "CONTATO"},
        {"owner": "RREZENDE", "name": "AMBIENTE"}
      ]
    }'
  );
END;
/

-- Profile for APP_SUPPORT
-- Context: support tickets, tasks and maintenance
BEGIN
  DBMS_CLOUD_AI.CREATE_PROFILE(
    profile_name => 'AI_PROFILE_SUPPORT',
    attributes   => '{
      "provider"          : "oci",
      "credential_name"   : "OCI$RESOURCE_PRINCIPAL",
      "model"             : "cohere.command-r-plus-08-2024",
      "oci_compartment_id": "ocid1.compartment.oc1..xxxxxxxx",
      "temperature"       : 0,
      "comments"          : true,
      "object_list"       : [
        {"owner": "RREZENDE", "name": "SUPORTE"},
        {"owner": "RREZENDE", "name": "TAREFA"},
        {"owner": "RREZENDE", "name": "MANUTENCAO"}
      ]
    }'
  );
END;
/

-- Profile for APP_HR
-- Context: people and task allocation
BEGIN
  DBMS_CLOUD_AI.CREATE_PROFILE(
    profile_name => 'AI_PROFILE_HR',
    attributes   => '{
      "provider"          : "oci",
      "credential_name"   : "OCI$RESOURCE_PRINCIPAL",
      "model"             : "cohere.command-r-plus-08-2024",
      "oci_compartment_id": "ocid1.compartment.oc1..xxxxxxxx",
      "temperature"       : 0,
      "comments"          : true,
      "object_list"       : [
        {"owner": "RREZENDE", "name": "CONTATO"},
        {"owner": "RREZENDE", "name": "TAREFA"}
      ]
    }'
  );
END;
/

-- Confirm profiles and registered tables
COL PROFILE_NAME    FOR A20
COL ATTRIBUTE_NAME  FOR A20
COL ATTRIBUTE_VALUE FOR A100

SELECT profile_name, attribute_name, attribute_value
  FROM user_cloud_ai_profile_attributes
 WHERE profile_name IN ('AI_PROFILE_SALES','AI_PROFILE_SUPPORT','AI_PROFILE_HR')
 ORDER BY profile_name, attribute_name;

GRANT SELECT ON rrezende.empresa    TO ai_ops;
GRANT SELECT ON rrezende.contato    TO ai_ops;
GRANT SELECT ON rrezende.ambiente   TO ai_ops;
GRANT SELECT ON rrezende.suporte    TO ai_ops;
GRANT SELECT ON rrezende.tarefa     TO ai_ops;
GRANT SELECT ON rrezende.manutencao TO ai_ops;

-- Create application schemas (if they don't exist yet)
CREATE USER app_sales   IDENTIFIED BY "&password"
  DEFAULT TABLESPACE data QUOTA UNLIMITED ON data;
CREATE USER app_hr      IDENTIFIED BY "&password"
  DEFAULT TABLESPACE data QUOTA UNLIMITED ON data;
CREATE USER app_support IDENTIFIED BY "&password"
  DEFAULT TABLESPACE data QUOTA UNLIMITED ON data;

GRANT CREATE SESSION TO app_sales, app_hr, app_support;

-- Grant EXECUTE on the gateway to each application schema
-- Do NOT grant EXECUTE on DBMS_CLOUD_AI directly
GRANT EXECUTE ON ai_ops.pkg_ai_gateway TO app_sales;
GRANT EXECUTE ON ai_ops.pkg_ai_gateway TO app_hr;
GRANT EXECUTE ON ai_ops.pkg_ai_gateway TO app_support;

-- Create a public synonym for convenience (optional)
CREATE PUBLIC SYNONYM pkg_ai_gateway FOR ai_ops.pkg_ai_gateway;

-- Connected as APP_SALES
DECLARE
  l_result CLOB;
BEGIN
  l_result := pkg_ai_gateway.generate(
    p_prompt       => 'Which companies have an active contract?',
    p_profile_name => 'AI_PROFILE_SALES',
    p_action       => 'runsql'
  );
  DBMS_OUTPUT.PUT_LINE(l_result);
END;
/

-- Or with narrate
DECLARE
  l_summary CLOB;
BEGIN
  l_summary := pkg_ai_gateway.generate(
    p_prompt       => 'Summarize companies by market segment',
    p_profile_name => 'AI_PROFILE_SALES',
    p_action       => 'narrate'
  );
  DBMS_OUTPUT.PUT_LINE(l_summary);
END;
/

-- Connected as APP_SUPPORT
DECLARE
  l_result CLOB;
BEGIN
  l_result := pkg_ai_gateway.generate(
    p_prompt       => 'Which support tickets have been open for more than 5 days?',
    p_profile_name => 'AI_PROFILE_SUPPORT',
    p_action       => 'runsql'
  );
  DBMS_OUTPUT.PUT_LINE(l_result);
END;
/

DECLARE
  l_summary CLOB;
BEGIN
  l_summary := pkg_ai_gateway.generate(
    p_prompt       => 'Summarize maintenance completed in the last month',
    p_profile_name => 'AI_PROFILE_SUPPORT',
    p_action       => 'narrate'
  );
  DBMS_OUTPUT.PUT_LINE(l_summary);
END;
/

-- Connected as APP_HR
DECLARE
  l_result CLOB;
BEGIN
  l_result := pkg_ai_gateway.generate(
    p_prompt       => 'Which contacts have tasks in progress?',
    p_profile_name => 'AI_PROFILE_HR',
    p_action       => 'runsql'
  );
  DBMS_OUTPUT.PUT_LINE(l_result);
END;
/

DECLARE
  l_summary CLOB;
BEGIN
  l_summary := pkg_ai_gateway.generate(
    p_prompt       => 'Summarize task distribution by contact in the last quarter',
    p_profile_name => 'AI_PROFILE_HR',
    p_action       => 'narrate'
  );
  DBMS_OUTPUT.PUT_LINE(l_summary);
END;
/

SELECT TRUNC(log_ts)           AS day_date,
       caller_schema,
       action,
       COUNT(*)                AS calls,
       SUM(tokens_input)       AS input_tokens,
       SUM(tokens_output)      AS output_tokens,
       SUM(tokens_total)       AS total_tokens,
       ROUND(AVG(duration_ms)) AS avg_ms,
       SUM(CASE WHEN status = 'ERROR' THEN 1 ELSE 0 END) AS errors
   FROM ai_ops.ai_token_log
   WHERE log_ts >= SYSDATE - 30
   GROUP BY TRUNC(log_ts), caller_schema, action
   ORDER BY day_date DESC, total_tokens DESC;

SELECT TO_CHAR(log_ts, 'YYYY-MM')  AS month,
       caller_schema,
       COUNT(*)                     AS total_calls,
       SUM(tokens_input)            AS input_tokens,
       SUM(tokens_output)           AS output_tokens,
       SUM(tokens_total)            AS total_tokens,
       ROUND(RATIO_TO_REPORT(SUM(tokens_total))
         OVER (PARTITION BY TO_CHAR(log_ts, 'YYYY-MM')) * 100, 2) AS pct_of_month
  FROM ai_ops.ai_token_log
  WHERE log_ts >= TRUNC(SYSDATE, 'YEAR')
    AND status = 'OK'
  GROUP BY TO_CHAR(log_ts, 'YYYY-MM'), caller_schema
  ORDER BY month DESC, total_tokens DESC;

SELECT *
  FROM (SELECT log_id,
               log_ts,
               caller_schema,
               profile_name,
               action,
               tokens_total,
               duration_ms,
               SUBSTR(prompt_text, 1, 200) AS prompt_preview
          FROM ai_ops.ai_token_log
          WHERE status = 'OK'
          ORDER BY tokens_total DESC)
  WHERE rownum <= 10;

COL month         FOR A10
COL caller_schema FOR A15
COL model_name    FOR A30
COL calls         FOR 999,999
COL total_chars   FOR 999,999,999
COL cost_brl      FOR 999,990.999999

-- Monthly cost by schema
-- SKU B108077: Large Cohere = BRL 0.09 per 10,000 transactions (hypothetical value)
-- 1 transaction = 1 character (prompt + response)
-- tokens * 4 converts estimated tokens back to characters
-- Update the unit price according to your current contract
SELECT TO_CHAR(log_ts, 'YYYY-MM')                       AS month,
       caller_schema,
       model_name,
       COUNT(*)                                          AS calls,
       SUM(tokens_total * 4)                             AS total_chars,
       ROUND(SUM(tokens_total * 4) / 10000 * 0.09, 6)   AS cost_brl
  FROM ai_ops.ai_token_log
  WHERE status  = 'OK'
    AND log_ts >= TRUNC(SYSDATE, 'YEAR')
  GROUP BY TO_CHAR(log_ts, 'YYYY-MM'), caller_schema, model_name
  ORDER BY month DESC, cost_brl DESC;

SELECT caller_schema,
       error_code,
       error_message,
       COUNT(*)    AS occurrences,
       MIN(log_ts) AS first_seen,
       MAX(log_ts) AS last_seen
  FROM ai_ops.ai_token_log
  WHERE status  = 'ERROR'
    AND log_ts >= ADD_MONTHS(SYSDATE, -1)
  GROUP BY caller_schema, error_code, error_message
  ORDER BY occurrences DESC;