Added adapter_source and api_token to completions API (#446)

predibase · Apr 29, 2024 · 16c3adc · 16c3adc
1 parent 6bc5f88
commit 16c3adc
Show file tree

Hide file tree

Showing 2 changed files with 114 additions and 39 deletions.
diff --git a/docs/reference/openapi.json b/docs/reference/openapi.json
@@ -969,6 +969,14 @@
  "stream": {
  "type": "boolean",
  "default": "false"
+ },
+ "adapter_source": {
+ "type": "string",
+ "nullable": true
+ },
+ "api_token": {
+ "type": "string",
+ "nullable": true
  }
  }
  },
@@ -1028,6 +1036,14 @@
  "stream": {
  "type": "boolean",
  "default": "false"
+ },
+ "adapter_source": {
+ "type": "string",
+ "nullable": true
+ },
+ "api_token": {
+ "type": "string",
+ "nullable": true
  }
  }
  },
@@ -1085,8 +1101,14 @@
  "ChatCompletionStreamResponseChoice": {
  "type": "object",
  "properties": {
- "index": { "type": "integer", "format": "int32" },
- "delta": { "type": "object", "$ref": "#/components/schemas/ChatCompletionResponseMessage" },
+ "index": {
+ "type": "integer",
+ "format": "int32"
+ },
+ "delta": {
+ "type": "object",
+ "$ref": "#/components/schemas/ChatCompletionResponseMessage"
+ },
  "finish_reason": {
  "type": "string",
  "nullable": true
@@ -1096,9 +1118,15 @@
  "ChatCompletionResponseChoice": {
  "type": "object",
  "properties": {
- "index": { "type": "integer", "format": "int32" },
- "message": { "type": "object", "$ref": "#/components/schemas/ChatCompletionResponseMessage" },
- "finish_reason": { 
+ "index": {
+ "type": "integer",
+ "format": "int32"
+ },
+ "message": {
+ "type": "object",
+ "$ref": "#/components/schemas/ChatCompletionResponseMessage"
+ },
+ "finish_reason": {
  "type": "string",
  "nullable": true
  }
@@ -1107,8 +1135,12 @@
  "ChatCompletionResponseMessage": {
  "type": "object",
  "properties": {
- "role": { "type": "string" },
- "content": { "type": "string" }
+ "role": {
+ "type": "string"
+ },
+ "content": {
+ "type": "string"
+ }
  }
  },
  "CompletionResponse": {
@@ -1141,14 +1173,19 @@
  "CompletionResponseChoice": {
  "type": "object",
  "properties": {
- "index": { "type": "integer", "format": "int32" },
- "text": { "type": "string" },
- "logprobs": { 
+ "index": {
+ "type": "integer",
+ "format": "int32"
+ },
+ "text": {
+ "type": "string"
+ },
+ "logprobs": {
  "type": "object",
  "nullable": true,
  "$ref": "#/components/schemas/LogProbs"
  },
- "finish_reason": { 
+ "finish_reason": {
  "type": "string",
  "nullable": true
  }
@@ -1157,15 +1194,26 @@
  "CompletionStreamResponse": {
  "type": "object",
  "properties": {
- "id": { "type": "string" },
- "object": { "type": "string" },
- "created": { "type": "integer", "format": "int64" },
- "model": { "type": "string" },
- "choices": { 
+ "id": {
+ "type": "string"
+ },
+ "object": {
+ "type": "string"
+ },
+ "created": {
+ "type": "integer",
+ "format": "int64"
+ },
+ "model": {
+ "type": "string"
+ },
+ "choices": {
  "type": "array",
- "items": { "$ref": "#/components/schemas/CompletionResponseStreamChoice" }
+ "items": {
+ "$ref": "#/components/schemas/CompletionResponseStreamChoice"
+ }
  },
- "usage": { 
+ "usage": {
  "type": "object",
  "nullable": true,
  "$ref": "#/components/schemas/UsageInfo"
@@ -1175,14 +1223,19 @@
  "CompletionResponseStreamChoice": {
  "type": "object",
  "properties": {
- "index": { "type": "integer", "format": "int32" },
- "text": { "type": "string" },
- "logprobs": { 
+ "index": {
+ "type": "integer",
+ "format": "int32"
+ },
+ "text": {
+ "type": "string"
+ },
+ "logprobs": {
  "type": "object",
  "nullable": true,
  "$ref": "#/components/schemas/LogProbs"
  },
- "finish_reason": { 
+ "finish_reason": {
  "type": "string",
  "nullable": true
  }
@@ -1191,24 +1244,36 @@
  "LogProbs": {
  "type": "object",
  "properties": {
- "text_offset": { 
+ "text_offset": {
  "type": "array",
- "items": { "type": "integer", "format": "int32" }
+ "items": {
+ "type": "integer",
+ "format": "int32"
+ }
  },
- "token_logprobs": { 
+ "token_logprobs": {
  "type": "array",
- "items": { "type": "number", "format": "float", "nullable": true }
+ "items": {
+ "type": "number",
+ "format": "float",
+ "nullable": true
+ }
  },
- "tokens": { 
+ "tokens": {
  "type": "array",
- "items": { "type": "string" }
+ "items": {
+ "type": "string"
+ }
  },
- "top_logprobs": { 
+ "top_logprobs": {
  "type": "array",
- "items": { 
+ "items": {
  "type": "object",
  "nullable": true,
- "additionalProperties": { "type": "number", "format": "float" }
+ "additionalProperties": {
+ "type": "number",
+ "format": "float"
+ }
  },
  "nullable": true
  }
@@ -1217,9 +1282,15 @@
  "UsageInfo": {
  "type": "object",
  "properties": {
- "prompt_tokens": { "type": "integer", "format": "uint32" },
- "total_tokens": { "type": "integer", "format": "uint32" },
- "completion_tokens": { 
+ "prompt_tokens": {
+ "type": "integer",
+ "format": "uint32"
+ },
+ "total_tokens": {
+ "type": "integer",
+ "format": "uint32"
+ },
+ "completion_tokens": {
  "type": "integer",
  "format": "uint32",
  "nullable": true
@@ -1453,4 +1524,4 @@
  "description": "LoRAX API"
  }
  ]
-}
+}
diff --git a/router/src/lib.rs b/router/src/lib.rs
@@ -502,6 +502,8 @@ struct ChatCompletionRequest {
  repetition_penalty: Option<f32>,
  top_k: Option<i32>,
  ignore_eos_token: Option<bool>,
+ adapter_source: Option<String>,
+ api_token: Option<String>,
 }
 
 #[derive(Clone, Debug, Deserialize, ToSchema)]
@@ -529,6 +531,8 @@ struct CompletionRequest {
  repetition_penalty: Option<f32>,
  top_k: Option<i32>,
  ignore_eos_token: Option<bool>,
+ adapter_source: Option<String>,
+ api_token: Option<String>,
 }
 
 #[derive(Serialize, ToSchema)]
@@ -635,9 +639,9 @@ impl From<CompletionRequest> for CompatGenerateRequest {
  inputs: req.prompt,
  parameters: GenerateParameters {
  adapter_id: req.model.parse().ok(),
- adapter_source: None,
+ adapter_source: req.adapter_source,
  adapter_parameters: None,
- api_token: None,
+ api_token: req.api_token,
  best_of: req.best_of.map(|x| x as usize),
  temperature: req.temperature,
  repetition_penalty: req.repetition_penalty,
@@ -669,9 +673,9 @@ impl From<ChatCompletionRequest> for CompatGenerateRequest {
  inputs: serde_json::to_string(&req.messages).unwrap(),
  parameters: GenerateParameters {
  adapter_id: req.model.parse().ok(),
- adapter_source: None,
+ adapter_source: req.adapter_source,
  adapter_parameters: None,
- api_token: None,
+ api_token: req.api_token,
  best_of: req.n.map(|x| x as usize),
  temperature: req.temperature,
  repetition_penalty: req.repetition_penalty,