Task: Redefine consistent output (WIP)

zya · Oct 9, 2023 · 90c7b76 · 90c7b76
1 parent 6dbf97b
commit 90c7b76
Show file tree

Hide file tree

Showing 5 changed files with 37 additions and 19 deletions.
diff --git a/src/handlers/anthropic.ts b/src/handlers/anthropic.ts
@@ -21,7 +21,6 @@ function toResponse(
   anthropicResponse: Anthropic.Completion,
 ): ResultNotStreaming {
   return {
-    id: '',
     choices: [
       {
         message: {
@@ -39,10 +38,9 @@ function toStreamingChunk(
   anthropicResponse: Anthropic.Completion,
 ): StreamingChunk {
   return {
-    id: '',
     choices: [
       {
-        delta: { content: anthropicResponse.completion },
+        delta: { content: anthropicResponse.completion, role: 'assistant' },
         finish_reason: 'stop',
         index: 0,
       },

diff --git a/src/handlers/cohere.ts b/src/handlers/cohere.ts
@@ -16,11 +16,11 @@ async function* toStream(
   response: cohereResponse<generateResponse>,
 ): AsyncIterable<StreamingChunk> {
   yield {
-    id: '',
     choices: [
       {
         delta: {
           content: response.body.generations[0].text,
+          role: 'assistant',
         },
         finish_reason: 'stop',
         index: 0,
@@ -61,7 +61,6 @@ export async function CohereHandler(
   }
 
   return {
-    id: '',
     choices: [
       {
         message: {

diff --git a/src/handlers/ollama.ts b/src/handlers/ollama.ts
@@ -17,10 +17,9 @@ interface OllamaResponseChunk {
 
 function toStreamingChunk(ollamaResponse: OllamaResponseChunk): StreamingChunk {
   return {
-    id: '',
     choices: [
       {
-        delta: { content: ollamaResponse.response },
+        delta: { content: ollamaResponse.response, role: 'assistant' },
         finish_reason: 'stop',
         index: 0,
       },
@@ -37,7 +36,6 @@ function toResponse(content: string): ResultNotStreaming {
         index: 0,
       },
     ],
-    id: '',
   };
 }
 

diff --git a/src/handlers/openai.ts b/src/handlers/openai.ts
@@ -24,5 +24,7 @@ export async function OpenAIHandler(
   params: HandlerParams,
 ): Promise<ResultNotStreaming | ResultStreaming> {
   const openai = new OpenAI();
-  return openai.chat.completions.create(params);
+  return openai.chat.completions.create(params) as Promise<
+    ResultNotStreaming | ResultStreaming
+  >; // TODO: Undo the type casting by properly handling and converting to consistent response
 }
diff --git a/src/types.ts b/src/types.ts
@@ -1,19 +1,40 @@
-import OpenAI from 'openai';
-
 export interface Message {
   role: 'system' | 'user' | 'assistant' | 'function';
   content: string | null;
 }
 
-export type ResultNotStreaming = Pick<
-  OpenAI.Chat.ChatCompletion,
-  'choices' | 'id'
->;
+export interface ConsistentResponseChoice {
+  finish_reason: string;
+  index: number;
+  message: {
+    role: string;
+    content: string;
+  };
+}
+
+export interface ConsistentResponseStreamingChoice
+  extends Omit<ConsistentResponseChoice, 'message'> {
+  delta: ConsistentResponseChoice['message'];
+}
+
+export interface ConsistentResponseUsage {
+  prompt_tokens: number;
+  completion_tokesn: number;
+  total_tokens: number;
+}
+
+export interface ConsistentResponse {
+  choices: ConsistentResponseChoice[];
+  model?: string; // TODO: Make this non-optional
+  created?: string; // TODO: Make this non-optional and implement
+  usage?: ConsistentResponseUsage; // TODO: Make this non-optional and implement
+}
 
-export type StreamingChunk = Pick<
-  OpenAI.Chat.ChatCompletionChunk,
-  'choices' | 'id'
->;
+export type ResultNotStreaming = ConsistentResponse;
+
+export interface StreamingChunk extends Omit<ConsistentResponse, 'choices'> {
+  choices: ConsistentResponseStreamingChoice[];
+}
 
 export type ResultStreaming = AsyncIterable<StreamingChunk>;