1
// Copyright (c) 2025 Cloudflare, Inc.
2
// Licensed under the Apache 2.0 license found in the LICENSE file or at:
3
//     https://opensource.org/licenses/Apache-2.0
4
 
5
interface Fetcher {
6
  fetch: typeof fetch;
7
}
8
 
9
export class AutoRAGInternalError extends Error {
10
  constructor(message: string, name = 'AutoRAGInternalError') {
11
    super(message);
12
    this.name = name;
13
  }
14
}
15
 
16
export class AutoRAGNotFoundError extends Error {
17
  constructor(message: string, name = 'AutoRAGNotFoundError') {
18
    super(message);
19
    this.name = name;
20
  }
21
}
22
 
23
export class AutoRAGUnauthorizedError extends Error {
24
  constructor(message: string, name = 'AutoRAGUnauthorizedError') {
25
    super(message);
26
    this.name = name;
27
  }
28
}
29
 
30
export class AutoRAGNameNotSetError extends Error {
31
  constructor(message: string, name = 'AutoRAGNameNotSetError') {
32
    super(message);
33
    this.name = name;
34
  }
35
}
36
 
37
async function parseError(
38
  res: Response,
39
  defaultMsg = 'Internal Error',
40
  errorCls = AutoRAGInternalError
41
): Promise<Error> {
42
  const content = await res.text();
43
 
44
  try {
45
    const parsedContent = JSON.parse(content) as {
46
      errors: { message: string }[];
47
    };
48
 
49
    return new errorCls(parsedContent.errors.at(0)?.message || defaultMsg);
50
  } catch {
51
    return new AutoRAGInternalError(content);
52
  }
53
}
54
 
55
export type ComparisonFilter = {
56
  key: string;
57
  type: 'eq' | 'ne' | 'gt' | 'gte' | 'lt' | 'lte';
58
  value: string | number | boolean;
59
};
60
 
61
export type CompoundFilter = {
62
  type: 'and' | 'or';
63
  filters: ComparisonFilter[];
64
};
65
 
66
export type AutoRagSearchRequest = {
67
  query: string;
68
  filters?: CompoundFilter | ComparisonFilter;
69
  max_num_results?: number;
70
  ranking_options?: {
71
    ranker?: string;
72
    score_threshold?: number;
73
  };
74
  reranking?: {
75
    enabled?: boolean;
76
    model?: string;
77
  };
78
  rewrite_query?: boolean;
79
};
80
 
81
export type AutoRagAiSearchRequest = AutoRagSearchRequest & {
82
  stream?: boolean;
83
  system_prompt?: string;
84
};
85
export type AutoRagAiSearchRequestStreaming = Omit<
86
  AutoRagAiSearchRequest,
87
  'stream'
88
> & {
89
  stream: true;
90
};
91
 
92
export type AutoRagSearchResponse = {
93
  object: 'vector_store.search_results.page';
94
  search_query: string;
95
  data: {
96
    file_id: string;
97
    filename: string;
98
    score: number;
99
    attributes: Record<string, string | number | boolean | null>;
100
    content: {
101
      type: 'text';
102
      text: string;
103
    }[];
104
  }[];
105
  has_more: boolean;
106
  next_page: string | null;
107
};
108
 
109
export type AutoRagListResponse = {
110
  id: string;
111
  enable: boolean;
112
  type: string;
113
  source: string;
114
  vectorize_name: string;
115
  paused: boolean;
116
  status: string;
117
}[];
118
 
119
export type AutoRagAiSearchResponse = AutoRagSearchResponse & {
120
  response: string;
121
};
122
 
123
export class AutoRAG {
124
  readonly #fetcher: Fetcher;
125
  readonly #autoragId: string | null;
126
 
127
  constructor(fetcher: Fetcher, autoragId?: string) {
128
    this.#fetcher = fetcher;
129
    this.#autoragId = autoragId || null;
130
  }
131
 
132
  async list(): Promise<AutoRagListResponse> {
133
    const res = await this.#fetcher.fetch(
134
      `https://workers-binding.ai/autorag/rags`,
135
      {
136
        method: 'GET',
137
        headers: {
138
          'content-type': 'application/json',
139
        },
140
      }
141
    );
142
 
143
    if (!res.ok) {
144
      throw await parseError(res);
145
    }
146
 
147
    const data = (await res.json()) as { result: AutoRagListResponse };
148
 
149
    return data.result;
150
  }
151
 
152
  async search(params: AutoRagSearchRequest): Promise<AutoRagSearchResponse> {
153
    if (!this.#autoragId) {
154
      throw new AutoRAGNameNotSetError('AutoRAG name not defined');
155
    }
156
 
157
    const res = await this.#fetcher.fetch(
158
      `https://workers-binding.ai/autorag/rags/${this.#autoragId}/search`,
159
      {
160
        method: 'POST',
161
        body: JSON.stringify(params),
162
        headers: {
163
          'content-type': 'application/json',
164
        },
165
      }
166
    );
167
 
168
    if (!res.ok) {
169
      if (res.status === 401) {
170
        throw await parseError(
171
          res,
172
          "You don't have access to this AutoRAG",
173
          AutoRAGUnauthorizedError
174
        );
175
      } else if (res.status === 404) {
176
        throw await parseError(res, 'AutoRAG not found', AutoRAGNotFoundError);
177
      }
178
      throw await parseError(res);
179
    }
180
 
181
    const data = (await res.json()) as { result: AutoRagSearchResponse };
182
 
183
    return data.result;
184
  }
185
 
186
  async aiSearch(params: AutoRagAiSearchRequestStreaming): Promise<Response>;
187
  async aiSearch(
188
    params: AutoRagAiSearchRequest
189
  ): Promise<AutoRagAiSearchResponse>;
190
  async aiSearch(
191
    params: AutoRagAiSearchRequest
192
  ): Promise<AutoRagAiSearchResponse | Response> {
193
    if (!this.#autoragId) {
194
      throw new AutoRAGNameNotSetError('AutoRAG name not defined');
195
    }
196
 
197
    const res = await this.#fetcher.fetch(
198
      `https://workers-binding.ai/autorag/rags/${this.#autoragId}/ai-search`,
199
      {
200
        method: 'POST',
201
        body: JSON.stringify(params),
202
        headers: {
203
          'content-type': 'application/json',
204
        },
205
      }
206
    );
207
 
208
    if (!res.ok) {
209
      if (res.status === 401) {
210
        throw await parseError(
211
          res,
212
          "You don't have access to this AutoRAG",
213
          AutoRAGUnauthorizedError
214
        );
215
      } else if (res.status === 404) {
216
        throw await parseError(res, 'AutoRAG not found', AutoRAGNotFoundError);
217
      }
218
      throw await parseError(res);
219
    }
220
 
221
    if (params.stream === true) {
222
      return res;
223
    }
224
 
225
    const data = (await res.json()) as { result: AutoRagAiSearchResponse };
226
 
227
    return data.result;
228
  }
229
}