Skip to content

Commit aa9feb3

Browse files
committed
web endpoints for web search and docs search
1 parent 1581297 commit aa9feb3

File tree

8 files changed

+747
-0
lines changed

8 files changed

+747
-0
lines changed

common/src/constants/analytics-events.ts

Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -100,6 +100,19 @@ export enum AnalyticsEvent {
100100
CHAT_COMPLETIONS_STREAM_STARTED = 'api.chat_completions_stream_started',
101101
CHAT_COMPLETIONS_ERROR = 'api.chat_completions_error',
102102

103+
// Web - Search API
104+
WEB_SEARCH_REQUEST = 'api.web_search_request',
105+
WEB_SEARCH_AUTH_ERROR = 'api.web_search_auth_error',
106+
WEB_SEARCH_VALIDATION_ERROR = 'api.web_search_validation_error',
107+
WEB_SEARCH_INSUFFICIENT_CREDITS = 'api.web_search_insufficient_credits',
108+
WEB_SEARCH_ERROR = 'api.web_search_error',
109+
110+
DOCS_SEARCH_REQUEST = 'api.docs_search_request',
111+
DOCS_SEARCH_AUTH_ERROR = 'api.docs_search_auth_error',
112+
DOCS_SEARCH_VALIDATION_ERROR = 'api.docs_search_validation_error',
113+
DOCS_SEARCH_INSUFFICIENT_CREDITS = 'api.docs_search_insufficient_credits',
114+
DOCS_SEARCH_ERROR = 'api.docs_search_error',
115+
103116
// Common
104117
FLUSH_FAILED = 'common.flush_failed',
105118
}

packages/agent-runtime/src/index.ts

Whitespace-only changes.
Lines changed: 143 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,143 @@
1+
import { afterEach, beforeEach, describe, expect, mock, test } from 'bun:test'
2+
import { NextRequest } from 'next/server'
3+
4+
import { postDocsSearch } from '../_post'
5+
6+
import type { TrackEventFn } from '@codebuff/common/types/contracts/analytics'
7+
import type {
8+
GetUserUsageDataFn,
9+
ConsumeCreditsWithFallbackFn,
10+
} from '@codebuff/common/types/contracts/billing'
11+
import type { GetUserInfoFromApiKeyFn } from '@codebuff/common/types/contracts/database'
12+
import type {
13+
Logger,
14+
LoggerWithContextFn,
15+
} from '@codebuff/common/types/contracts/logger'
16+
17+
describe('/api/v1/docs-search POST endpoint', () => {
18+
let mockLogger: Logger
19+
let mockLoggerWithContext: LoggerWithContextFn
20+
let mockTrackEvent: TrackEventFn
21+
let mockGetUserUsageData: GetUserUsageDataFn
22+
let mockGetUserInfoFromApiKey: GetUserInfoFromApiKeyFn
23+
let mockConsumeCreditsWithFallback: ConsumeCreditsWithFallbackFn
24+
let mockFetch: typeof globalThis.fetch
25+
26+
beforeEach(() => {
27+
mockLogger = {
28+
error: mock(() => {}),
29+
warn: mock(() => {}),
30+
info: mock(() => {}),
31+
debug: mock(() => {}),
32+
}
33+
mockLoggerWithContext = mock(() => mockLogger)
34+
mockTrackEvent = mock(() => {})
35+
36+
mockGetUserUsageData = mock(async () => ({
37+
balance: { totalRemaining: 10 },
38+
nextQuotaReset: 'soon',
39+
}))
40+
mockGetUserInfoFromApiKey = mock(async ({ apiKey }) =>
41+
apiKey === 'valid' ? ({ id: 'user-1' } as any) : null,
42+
)
43+
mockConsumeCreditsWithFallback = mock(
44+
async () =>
45+
({ success: true, value: { chargedToOrganization: false } }) as any,
46+
)
47+
48+
// Mock fetch for Context7 search and docs endpoints
49+
mockFetch = (async (url: any) => {
50+
const u = typeof url === 'string' ? new URL(url) : url
51+
if (String(u).includes('/search')) {
52+
return new Response(
53+
JSON.stringify({
54+
results: [
55+
{
56+
id: 'lib1',
57+
title: 'Lib1',
58+
description: '',
59+
branch: 'main',
60+
lastUpdateDate: '',
61+
state: 'finalized',
62+
totalTokens: 100,
63+
totalSnippets: 10,
64+
totalPages: 1,
65+
},
66+
],
67+
}),
68+
{ status: 200, headers: { 'Content-Type': 'application/json' } },
69+
)
70+
}
71+
return new Response('Some documentation text', {
72+
status: 200,
73+
headers: { 'Content-Type': 'text/plain' },
74+
})
75+
}) as any
76+
})
77+
78+
afterEach(() => {
79+
mock.restore()
80+
})
81+
82+
test('401 when missing API key', async () => {
83+
const req = new NextRequest('http://localhost:3000/api/v1/docs-search', {
84+
method: 'POST',
85+
body: JSON.stringify({ libraryTitle: 'React' }),
86+
})
87+
const res = await postDocsSearch({
88+
req,
89+
getUserInfoFromApiKey: mockGetUserInfoFromApiKey,
90+
logger: mockLogger,
91+
loggerWithContext: mockLoggerWithContext,
92+
trackEvent: mockTrackEvent,
93+
getUserUsageData: mockGetUserUsageData,
94+
consumeCreditsWithFallback: mockConsumeCreditsWithFallback,
95+
fetch: mockFetch,
96+
})
97+
expect(res.status).toBe(401)
98+
})
99+
100+
test('402 when insufficient credits', async () => {
101+
mockGetUserUsageData = mock(async () => ({
102+
balance: { totalRemaining: 0 },
103+
nextQuotaReset: 'soon',
104+
}))
105+
const req = new NextRequest('http://localhost:3000/api/v1/docs-search', {
106+
method: 'POST',
107+
headers: { Authorization: 'Bearer valid' },
108+
body: JSON.stringify({ libraryTitle: 'React' }),
109+
})
110+
const res = await postDocsSearch({
111+
req,
112+
getUserInfoFromApiKey: mockGetUserInfoFromApiKey,
113+
logger: mockLogger,
114+
loggerWithContext: mockLoggerWithContext,
115+
trackEvent: mockTrackEvent,
116+
getUserUsageData: mockGetUserUsageData,
117+
consumeCreditsWithFallback: mockConsumeCreditsWithFallback,
118+
fetch: mockFetch,
119+
})
120+
expect(res.status).toBe(402)
121+
})
122+
123+
test('200 on success', async () => {
124+
const req = new NextRequest('http://localhost:3000/api/v1/docs-search', {
125+
method: 'POST',
126+
headers: { Authorization: 'Bearer valid' },
127+
body: JSON.stringify({ libraryTitle: 'React', topic: 'Hooks' }),
128+
})
129+
const res = await postDocsSearch({
130+
req,
131+
getUserInfoFromApiKey: mockGetUserInfoFromApiKey,
132+
logger: mockLogger,
133+
loggerWithContext: mockLoggerWithContext,
134+
trackEvent: mockTrackEvent,
135+
getUserUsageData: mockGetUserUsageData,
136+
consumeCreditsWithFallback: mockConsumeCreditsWithFallback,
137+
fetch: mockFetch,
138+
})
139+
expect(res.status).toBe(200)
140+
const body = await res.json()
141+
expect(body.documentation).toContain('Some documentation text')
142+
})
143+
})
Lines changed: 217 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,217 @@
1+
import { AnalyticsEvent } from '@codebuff/common/constants/analytics-events'
2+
import { PROFIT_MARGIN } from '@codebuff/common/old-constants'
3+
import { NextResponse } from 'next/server'
4+
import { z } from 'zod'
5+
6+
import type { TrackEventFn } from '@codebuff/common/types/contracts/analytics'
7+
import type {
8+
GetUserUsageDataFn,
9+
ConsumeCreditsWithFallbackFn,
10+
} from '@codebuff/common/types/contracts/billing'
11+
import type { GetUserInfoFromApiKeyFn } from '@codebuff/common/types/contracts/database'
12+
import type {
13+
Logger,
14+
LoggerWithContextFn,
15+
} from '@codebuff/common/types/contracts/logger'
16+
import type { NextRequest } from 'next/server'
17+
18+
import { fetchContext7LibraryDocumentation } from '@codebuff/agent-runtime/llm-api/context7-api'
19+
import { extractApiKeyFromHeader } from '@/util/auth'
20+
21+
const bodySchema = z.object({
22+
libraryTitle: z.string().min(1, 'libraryTitle is required'),
23+
topic: z.string().optional(),
24+
maxTokens: z.number().int().positive().optional(),
25+
repoUrl: z.string().url().optional(),
26+
})
27+
28+
export async function postDocsSearch(params: {
29+
req: NextRequest
30+
getUserInfoFromApiKey: GetUserInfoFromApiKeyFn
31+
logger: Logger
32+
loggerWithContext: LoggerWithContextFn
33+
trackEvent: TrackEventFn
34+
getUserUsageData: GetUserUsageDataFn
35+
consumeCreditsWithFallback: ConsumeCreditsWithFallbackFn
36+
fetch: typeof globalThis.fetch
37+
}) {
38+
const {
39+
req,
40+
getUserInfoFromApiKey,
41+
loggerWithContext,
42+
trackEvent,
43+
getUserUsageData,
44+
consumeCreditsWithFallback,
45+
fetch,
46+
} = params
47+
let { logger } = params
48+
49+
// Parse JSON body
50+
let json: unknown
51+
try {
52+
json = await req.json()
53+
} catch (e) {
54+
trackEvent({
55+
event: AnalyticsEvent.DOCS_SEARCH_VALIDATION_ERROR,
56+
userId: 'unknown',
57+
properties: { error: 'Invalid JSON' },
58+
logger,
59+
})
60+
return NextResponse.json(
61+
{ error: 'Invalid JSON in request body' },
62+
{ status: 400 },
63+
)
64+
}
65+
66+
// Validate body
67+
const parsed = bodySchema.safeParse(json)
68+
if (!parsed.success) {
69+
trackEvent({
70+
event: AnalyticsEvent.DOCS_SEARCH_VALIDATION_ERROR,
71+
userId: 'unknown',
72+
properties: { issues: parsed.error.format() },
73+
logger,
74+
})
75+
return NextResponse.json(
76+
{ error: 'Invalid request body', details: parsed.error.format() },
77+
{ status: 400 },
78+
)
79+
}
80+
const { libraryTitle, topic, maxTokens, repoUrl } = parsed.data
81+
82+
// Auth
83+
const apiKey = extractApiKeyFromHeader(req)
84+
if (!apiKey) {
85+
trackEvent({
86+
event: AnalyticsEvent.DOCS_SEARCH_AUTH_ERROR,
87+
userId: 'unknown',
88+
properties: { reason: 'Missing API key' },
89+
logger,
90+
})
91+
return NextResponse.json({ message: 'Unauthorized' }, { status: 401 })
92+
}
93+
94+
const userInfo = await getUserInfoFromApiKey({
95+
apiKey,
96+
fields: ['id', 'email', 'discord_id'],
97+
logger,
98+
})
99+
if (!userInfo) {
100+
trackEvent({
101+
event: AnalyticsEvent.DOCS_SEARCH_AUTH_ERROR,
102+
userId: 'unknown',
103+
properties: { reason: 'Invalid API key' },
104+
logger,
105+
})
106+
return NextResponse.json(
107+
{ message: 'Invalid Codebuff API key' },
108+
{ status: 401 },
109+
)
110+
}
111+
logger = loggerWithContext({ userInfo })
112+
const userId = userInfo.id
113+
114+
// Track request
115+
trackEvent({
116+
event: AnalyticsEvent.DOCS_SEARCH_REQUEST,
117+
userId,
118+
properties: { libraryTitle, hasTopic: !!topic, hasRepoUrl: !!repoUrl },
119+
logger,
120+
})
121+
122+
// Credit cost: flat 1 credit (+profit margin)
123+
const baseCost = 1
124+
const creditsToCharge = Math.round(baseCost * (1 + PROFIT_MARGIN))
125+
126+
// Check credits
127+
const {
128+
balance: { totalRemaining },
129+
nextQuotaReset,
130+
} = await getUserUsageData({ userId, logger })
131+
if (totalRemaining <= 0 || totalRemaining < creditsToCharge) {
132+
trackEvent({
133+
event: AnalyticsEvent.DOCS_SEARCH_INSUFFICIENT_CREDITS,
134+
userId,
135+
properties: { totalRemaining, required: creditsToCharge, nextQuotaReset },
136+
logger,
137+
})
138+
return NextResponse.json(
139+
{
140+
message: 'Insufficient credits',
141+
totalRemaining,
142+
required: creditsToCharge,
143+
nextQuotaReset,
144+
},
145+
{ status: 402 },
146+
)
147+
}
148+
149+
// Charge upfront with delegation fallback
150+
const chargeResult = await consumeCreditsWithFallback({
151+
userId,
152+
creditsToCharge,
153+
repoUrl,
154+
context: 'documentation lookup',
155+
logger,
156+
})
157+
if (!chargeResult.success) {
158+
logger.error(
159+
{ userId, creditsToCharge, error: chargeResult.error },
160+
'Failed to charge credits for docs search',
161+
)
162+
return NextResponse.json(
163+
{ error: 'Failed to charge credits' },
164+
{ status: 500 },
165+
)
166+
}
167+
168+
// Perform docs fetch
169+
try {
170+
const documentation = await fetchContext7LibraryDocumentation({
171+
query: libraryTitle,
172+
topic,
173+
tokens: maxTokens,
174+
logger,
175+
fetch,
176+
})
177+
178+
if (!documentation) {
179+
trackEvent({
180+
event: AnalyticsEvent.DOCS_SEARCH_ERROR,
181+
userId,
182+
properties: { reason: 'No documentation' },
183+
logger,
184+
})
185+
return NextResponse.json(
186+
{
187+
error: `No documentation found for "${libraryTitle}"${topic ? ` with topic "${topic}"` : ''}`,
188+
},
189+
{ status: 200 },
190+
)
191+
}
192+
193+
return NextResponse.json({ documentation, creditsUsed: creditsToCharge })
194+
} catch (error) {
195+
logger.error(
196+
{
197+
error:
198+
error instanceof Error
199+
? { name: error.name, message: error.message, stack: error.stack }
200+
: error,
201+
},
202+
'Docs search failed',
203+
)
204+
trackEvent({
205+
event: AnalyticsEvent.DOCS_SEARCH_ERROR,
206+
userId,
207+
properties: {
208+
error: error instanceof Error ? error.message : 'Unknown error',
209+
},
210+
logger,
211+
})
212+
return NextResponse.json(
213+
{ error: 'Error fetching documentation' },
214+
{ status: 500 },
215+
)
216+
}
217+
}

0 commit comments

Comments
 (0)