1import * as _ from 'lodash-es'
2import { assertOnlyItem } from '../../sharedF/assertsF/assertOnlyItemF.mts'
3import { betLog } from '../../sharedF/betterConsLogF.mts'
4import { calcHash } from '../calcHashF.mts'
5import { getPlainMappedSec } from '../secretsF/getReqJsonSecretHF.mts'
6import { tsSecs } from '../secretsF/tsSecF.mts'
7import { json1Parse } from '../../sharedF/jsonF.mts'
8import { luxNow, luxUtcNow } from '../../sharedF/timeF/luxon0NowF.mts'
9import { runMemoTempfile } from '../ptDiskF/memoTmpfileF.mts'
10import { throPtErr } from '../../sharedF/throwErrorF/sharedErrF.mts'
11import { assertDefined } from '../../sharedF/assertsF/assertDefinedF.mts'
12import { assertNonEmpty } from '../../sharedF/assertsF/assertNonEmptyF.mts'
13import { throwIf } from '../../sharedF/throwErrorF/throwIfF.mts'
14import { markModelFailed } from './markModelFailedAI.mts'
15import { genContext } from '../genContextF.mts'
16
17export const aiAbortSig = new AbortController()
18export const aiCoderModel = 'anthropic/claude-4-opus-20250522'
19
20export const getOpenRouterApiToken = () => getPlainMappedSec(tsSecs.openrouterApiToken)
21
22type OpnRtrApiOpts = {method: string, bodyH?: Record<string, unknown>, apiKey?: string}
23
24export const opnRtrApi = async (urlPath: string, {method, bodyH, apiKey}: OpnRtrApiOpts) => {
25  const body = bodyH ? JSON.stringify(bodyH) : undefined
26  const token = apiKey || getOpenRouterApiToken()
27  const headers: Record<string, string> = {
28    'Authorization': `Bearer ${token}`,
29    'Content-Type': 'application/json',
30  }
31  if (bodyH) {
32    Object.assign(headers, {
33      'Idempotency-Key': calcHash(body!),
34    })
35  }
36  const ret = await fetch('https://openrouter.ai' + urlPath, {
37    signal: aiAbortSig.signal,
38    method,
39    headers,
40    body
41  })
42
43  return ret
44}
45
46type AiMessage = {role: string, content: string | unknown[]}
47
48export type AiUsageEntry = {label: string, model: string, promptTokens: number, completionTokens: number, costUsd: number | null}
49export const aiUsageCtx = genContext<AiUsageEntry[]>('aiUsage')
50export const recordAiUsage = ({label, model, usage}: {label: string, model: string, usage: any}) => {
51  const store = aiUsageCtx.getStore()
52  if (!store || !usage) return
53  store.push({label, model, promptTokens: usage.prompt_tokens ?? 0, completionTokens: usage.completion_tokens ?? 0, costUsd: usage.cost ?? null})
54}
55
56export const costUsdForUsage = async ({model, usage}: {model: string, usage: {cost?: number | null, prompt_tokens: number, completion_tokens: number}}): Promise<number | null> => {
57  if (usage.cost != null) return usage.cost
58  try {
59    const {pricing} = await getModelInfo({model})
60    return usage.prompt_tokens * parseFloat(pricing.prompt) + usage.completion_tokens * parseFloat(pricing.completion)
61  } catch { return null }
62}
63
64export const openRouterPrompt = async ({promptStr, aiMessages: messages, stream=false, model, apiKey, baseUrl = 'https://openrouter.ai/api/v1', trackCost = true, ...paramsH}: {promptStr?: string, aiMessages?: AiMessage[], stream?: boolean, model: string, apiKey?: string, baseUrl?: string, trackCost?: boolean, [k: string]: unknown}) => {
65  messages ||= [
66    {
67      role: 'user',
68      content: promptStr!,
69    }
70  ]
71
72  const startDt = luxUtcNow()
73  const bodyH = {
74    ...(trackCost ? { require_parameters: true } : {}),
75    ...paramsH,
76    stream, model,
77    messages
78  }
79  const body = JSON.stringify(bodyH)
80  const token = apiKey || getOpenRouterApiToken()
81  const resp = await fetch(`${baseUrl}/chat/completions`, {
82    signal: aiAbortSig.signal,
83    method: 'POST',
84    headers: {
85      'Authorization': `Bearer ${token}`,
86      'Content-Type': 'application/json',
87      'Idempotency-Key': calcHash(body),
88    },
89    body,
90  })
91
92  if (stream) {
93    let totContent = ''
94    for await (const chunk of resp.body as any) {
95      const str = new TextDecoder().decode(chunk)
96      const match = str.match(/data\: (.+)/)
97      if (match) {
98        const obj = json1Parse(match[1]) as any
99        const {content} = obj.choices[0].delta
100        totContent += content
101        process.stdout.write(content)
102      } else {
103        // console.log('streamed', str)
104      }
105    }
106    process.stdout.write('\n')
107    return totContent
108  } else {
109    const respH: any = await resp.json()
110
111    const {error} = respH
112    if (error) {
113      const isModelErr = error.code === 400 || /model|invalid/i.test(error.message || '')
114      if (isModelErr) markModelFailed(model)
115      throPtErr('openRouterErr', {error})
116    }
117
118    if (trackCost) {
119      const {usage} = respH
120      assertDefined(usage, {respH})
121      const costUsd = await costUsdForUsage({model, usage})
122      assertDefined(costUsd, {model, usage})
123
124      betLog({
125        dur: luxNow().diff(startDt),
126        prompt_tokens: usage.prompt_tokens,
127        completion_tokens: usage.completion_tokens,
128        costCents: costUsd * 100,
129      })
130
131      respH.costUsd = costUsd
132    } else {
133      respH.costUsd = respH.usage?.cost ?? 0
134    }
135    return respH
136  }
137}
138
139export const aiUsrMsg = (content: string) => {
140  assertDefined(content)
141  return {role: 'user', content}
142}
143
144export const aiAssMsg = (content: string) => {
145  assertDefined(content)
146  return {role: 'assistant', content}
147}
148
149export const opnRtrStr = async (props: {promptStr?: string, aiMessages?: AiMessage[], model: string, apiKey?: string, [k: string]: unknown}) => {
150  return (await openRouterPrompt(props)).choices[0].message.content
151}
152
153export const opnRtrStrSimple = async ({ model, promptStr, apiKey, max_tokens = 1024, baseUrl = 'https://openrouter.ai/api/v1', label = 'OpenRouter' }: {model: string, promptStr: string, apiKey: string, max_tokens?: number, baseUrl?: string, label?: string}) => {
154  const resp = await fetch(`${baseUrl}/chat/completions`, {
155    method: 'POST',
156    headers: {
157      'Authorization': `Bearer ${apiKey}`,
158      'Content-Type': 'application/json',
159    },
160    body: JSON.stringify({
161      model,
162      max_tokens,
163      messages: [{ role: 'user', content: promptStr }],
164    }),
165  })
166  if (!resp.ok) {
167    const errBody = await resp.text().catch(() => '')
168    if (resp.status === 400 || resp.status === 404 || /model|invalid/i.test(errBody)) markModelFailed(model)
169    throw new Error(`${label} ${resp.status}: ${errBody.slice(0, 200)}`)
170  }
171  const data: any = await resp.json()
172  if (data.error) {
173    if (data.error.code === 400 || /model|invalid/i.test(data.error.message || '')) markModelFailed(model)
174    throw new Error(`${label}: ${data.error.message || JSON.stringify(data.error)}`)
175  }
176  recordAiUsage({label, model, usage: data.usage})
177  return data.choices[0].message.content
178}
179
180export const getAllOpenRouterModelIds = async (): Promise<string[]> => {
181  return await runMemoTempfile({cacheKeyA: ['openrouter-all-models-v1']}, async () => {
182    const resp = await opnRtrApi(`/api/v1/models`, {method: 'GET'})
183    const data: any = await resp.json()
184    assertNonEmpty(data.data, {data})
185    return _.map(data.data, 'id') as string[]
186  })
187}
188
189export const getModelInfo = async ({model}: {model: string}): Promise<any> => {
190  return await runMemoTempfile({cacheKeyA: [model]}, async () => {
191    const allModelIds = await getAllOpenRouterModelIds()
192    const resp = await opnRtrApi(`/api/v1/models`, {method: 'GET'})
193    const data: any = await resp.json()
194    const matched = _.filter(data.data, ({id, canonical_slug}: {id: string, canonical_slug: string}) => {
195      return id == model || canonical_slug == model
196    })
197    assertNonEmpty(matched, {model, allModelIds: allModelIds.length})
198    return assertOnlyItem(matched)
199  })
200}
201