Update apis/chat_api.py

apis/chat_api.py

@@ -146,6 +146,73 @@ class ChatAPIApp:
             raise HTTPException(status_code=e.status_code, detail=e.detail)
         except Exception as e:
             raise HTTPException(status_code=500, detail=str(e))
+
+
+    class ChatCompletionsPostItem(BaseModel):
+        model: str = Field(
+            default="nous-mixtral-8x7b",
+            description="(str) `nous-mixtral-8x7b`",
+        )
+        messages: list = Field(
+            default=[{"role": "user", "content": "Hello, who are you?"}],
+            description="(list) Messages",
+        )
+        temperature: Union[float, None] = Field(
+            default=0.5,
+            description="(float) Temperature",
+        )
+        top_p: Union[float, None] = Field(
+            default=0.95,
+            description="(float) top p",
+        )
+        max_tokens: Union[int, None] = Field(
+            default=-1,
+            description="(int) Max tokens",
+        )
+        use_cache: bool = Field(
+            default=False,
+            description="(bool) Use cache",
+        )
+        stream: bool = Field(
+            default=True,
+            description="(bool) Stream",
+        )
+
+    def chat_completions_ollama(
+        self, item: ChatCompletionsPostItem, api_key: str = Depends(extract_api_key)
+    ):
+        try:
+            print(item.messages)
+            item.model = "llama3-8b" if item.model == "llama3" else item.model
+            api_key = self.auth_api_key(api_key)
+
+            if item.model == "gpt-3.5-turbo":
+                streamer = OpenaiStreamer()
+                stream_response = streamer.chat_response(messages=item.messages)
+            elif item.model in PRO_MODELS:
+                streamer = HuggingchatStreamer(model=item.model)
+                stream_response = streamer.chat_response(
+                    messages=item.messages,
+                )
+            else:
+                streamer = HuggingfaceStreamer(model=item.model)
+                composer = MessageComposer(model=item.model)
+                composer.merge(messages=item.messages)
+                stream_response = streamer.chat_response(
+                    prompt=composer.merged_str,
+                    temperature=item.temperature,
+                    top_p=item.top_p,
+                    max_new_tokens=item.max_tokens,
+                    api_key=api_key,
+                    use_cache=item.use_cache,
+                )
+
+            data_response = streamer.chat_return_dict(stream_response)
+            return data_response
+        except HfApiException as e:
+            raise HTTPException(status_code=e.status_code, detail=e.detail)
+        except Exception as e:
+            raise HTTPException(status_code=500, detail=str(e))
             
 
     class GenerateRequest(BaseModel):
@@ -315,7 +382,7 @@ class ChatAPIApp:
                 prefix + "/chat",
                 summary="Ollama Chat completions in conversation session",
                 include_in_schema=include_in_schema,
-            )(self.chat_completions)
+            )(self.chat_completions_ollama)
             
             self.app.post(
                 prefix + "/embeddings",