feat: convrag added to client, fix response mapping

Author: amirai21

package.json

@@ -20,7 +20,8 @@
     "lint": "npx eslint 'src/**/*.{ts,tsx}' --no-ignore",
     "format": "prettier --write \"(src|test)/**\" --no-error-on-unmatched-pattern",
     "prepare": "npm run build",
-    "example": "npx tsx src/example.ts",
+    "chatExample": "npx tsx src/examples/chatExample.ts",
+    "convRagExample": "npx tsx src/examples/convRagExample.ts",
     "circular": "madge --circular --extensions ts src",
     "quality": "npm run circular && npm run lint && tsc --noEmit && npm run format && npm run unused-deps",
     "quality:fix": "npm run circular && npm run lint -- --fix && tsc --noEmit && npm run format"

src/AI21.ts

@@ -4,6 +4,7 @@ import { MissingAPIKeyError } from './errors';
 import { Chat } from './resources/chat';
 import { APIClient } from './APIClient';
 import { Headers } from './types';
+import { ConversationalRag } from './resources/rag/conversationalRag';
 
 export type ClientOptions = {
   baseURL?: string;
@@ -54,6 +55,7 @@ export class AI21 extends APIClient {
 
   // Resources
   chat: Chat = new Chat(this);
+  conversationalRag: ConversationalRag = new ConversationalRag(this);
 
   // eslint-disable-next-line @typescript-eslint/no-unused-vars
   protected override authHeaders(_: Types.FinalRequestOptions): Types.Headers {

src/ResponseHandler.ts

@@ -11,7 +11,7 @@ type APIResponse<T> = {
 export async function handleAPIResponse<T>({
   response,
   options,
-}: APIResponseProps): Promise<Stream<T> | APIResponse<T>> {
+}: APIResponseProps): Promise<Stream<T> | Promise<APIResponse<T>>> {
   if (options.stream) {
     if (!response.body) {
       throw new AI21Error('Response body is null');
@@ -20,10 +20,5 @@ export async function handleAPIResponse<T>({
   }
 
   const contentType = response.headers.get('content-type');
-  const data = contentType?.includes('application/json') ? await response.json() : null;
-
-  return {
-    data,
-    response,
-  };
+  return contentType?.includes('application/json') ? await response.json() : null;
 }

src/example.ts renamed to src/examples/chatExample.ts

@@ -1,19 +1,32 @@
-import { AI21 } from './AI21';
+import { AI21 } from "../AI21";
+
 /*
   This is a temporary example to test the API streaming/non-streaming functionality.
 */
 async function main() {
   const client = new AI21({ apiKey: process.env.AI21_API_KEY });
 
   try {
-    const response = await client.chat.completions.create({
+    
+    console.log('-------------------------------- streaming flow --------------------------------');
+
+    let streamResponse = await client.chat.completions.create({
       model: 'jamba-1.5-mini',
       messages: [{ role: 'user', content: 'Hello, how are you? tell me a 100 line story about a cat' }],
       stream: true,
     });
-    for await (const chunk of response) {
+    for await (const chunk of streamResponse) {
       process.stdout.write(chunk.choices[0]?.delta?.content || '');
     }
+
+    console.log('-------------------------------- non streaming flow --------------------------------');
+
+    const response = await client.chat.completions.create({
+      model: 'jamba-1.5-mini',
+      messages: [{ role: 'user', content: 'Hello, how are you? tell me a 100 line story about a cat' }],
+      stream: false,
+    });
+    console.log(response);
   } catch (error) {
     console.error('Error:', error);
   }

src/examples/convRagExample.ts

@@ -0,0 +1,31 @@
+import { AI21 } from "../AI21";
+
+/*
+  This is a temporary example to test the Conversational RAG functionality.
+*/
+async function main() {
+
+    /* TODO - add a file upload example when library support is added and combined with the below flow */
+
+  const client = new AI21({ apiKey: process.env.AI21_API_KEY });
+  try {
+
+    /* The following example is for a question that is not in the context of files uploaded to RAG */
+    
+    const answer_not_in_ctx_response = await client.conversationalRag.create({
+      messages: [{ role: 'user', content: 'Who is the Russian president?' }],
+    });
+    console.log(answer_not_in_ctx_response);
+
+    /* The following example is for a question that should be answered based on files uploaded to RAG */
+
+    const answer_in_ctx_response = await client.conversationalRag.create({
+      messages: [{ role: 'user', content: 'What is headace?' }],
+    });
+    console.log(answer_in_ctx_response);
+  } catch (error) {
+    console.error('Error:', error);
+  }
+}
+
+main().catch(console.error);

src/resources/rag/conversationalRag.ts

@@ -0,0 +1,17 @@
+import * as Models from '../../types';
+import { APIResource } from '../../APIResource';
+import { ConversationalRagRequest } from '../../types/rag/ConversationalRagRequest';
+import { ConversationalRagResponse } from '../../types/rag/ConversationalRagResponse';
+
+export class ConversationalRag extends APIResource {
+
+  create(body: ConversationalRagRequest, options?: Models.RequestOptions){
+    return this.client.post<ConversationalRagRequest, ConversationalRagResponse>(
+      '/conversational-rag',
+      {
+        body,
+        ...options,
+      } as Models.RequestOptions<ConversationalRagRequest>,
+    ) as Promise<ConversationalRagResponse>;
+  }
+}

src/types/rag/ConversationalRagRequest.ts

@@ -0,0 +1,15 @@
+import { ChatMessageParam } from "../chat";
+import { RetrievalStrategy } from "./RetrievalStrategy";
+
+export interface ConversationalRagRequest {
+    messages: ChatMessageParam[];
+    path?: string | null;
+    labels?: string[] | null;
+    file_ids?: string[] | null;
+    max_segments?: number | null;
+    retrieval_strategy?: RetrievalStrategy | string | null;
+    retrieval_similarity_threshold?: number | null;
+    max_neighbors?: number | null;
+    hybrid_search_alpha?: number | null;
+    [key: string]: any;
+  }

src/types/rag/ConversationalRagResponse.ts

@@ -0,0 +1,11 @@
+import { ChatMessage } from "../chat";
+import { ConversationalRagSource } from "./ConversationalRagSource";
+
+export interface ConversationalRagResponse {
+    id: string;
+    choices: ChatMessage[];
+    search_queries?: string[] | null;
+    context_retrieved: boolean;
+    answer_in_context: boolean;
+    sources: ConversationalRagSource[];
+}

src/types/rag/ConversationalRagSource.ts

@@ -0,0 +1,9 @@
+export interface ConversationalRagSource {
+    text: string;
+    file_id: string;
+    file_name: string;
+    score: number;
+    order?: number | null;
+    public_url?: string | null;
+    labels?: string[] | null;
+}

src/types/rag/RetrievalStrategy.ts

@@ -0,0 +1,6 @@
+export enum RetrievalStrategy {
+    DEFAULT = "default",
+    SEGMENTS = "segments",
+    ADD_NEIGHBORS = "add_neighbors",
+    FULL_DOC = "full_doc"
+}