Update app.py

app.py

@@ -24,58 +24,44 @@
 from fastapi import FastAPI, HTTPException
 from pydantic import BaseModel
 from huggingface_hub import InferenceClient
-from typing import List, Tuple
 
-# Initialize the Inference client
-client = InferenceClient("HuggingFaceH4/zephyr-7b-beta")
-
-# Initialize FastAPI
 app = FastAPI()
 
-class Message(BaseModel):
-    message: str
-    history: List[Tuple[str, str]]
-    system_message: str
-    max_tokens: int
-    temperature: float
-    top_p: float
-
-def generate_response(message, history: list[tuple[str, str]], system_message, max_tokens, temperature, top_p):
-    messages = [{"role": "system", "content": system_message}]
+client = InferenceClient("HuggingFaceH4/zephyr-7b-beta")
 
-    for val in history:
-        if val[0]:
-            messages.append({"role": "user", "content": val[0]})
-        if val[1]:
-            messages.append({"role": "assistant", "content": val[1]})
+class ChatRequest(BaseModel):
+    message: str
+    history: list[tuple[str, str]] = []
+    system_message: str = "You are a friendly Chatbot."
+    max_tokens: int = 512
+    temperature: float = 0.7
+    top_p: float = 0.95
 
-    messages.append({"role": "user", "content": message})
+class ChatResponse(BaseModel):
+    response: str
 
-    response = ""
+@app.post("/chat", response_model=ChatResponse)
+async def chat(request: ChatRequest):
+    try:
+        messages = [{"role": "system", "content": request.system_message}]
+        for val in request.history:
+            if val[0]:
+                messages.append({"role": "user", "content": val[0]})
+            if val[1]:
+                messages.append({"role": "assistant", "content": val[1]})
+        messages.append({"role": "user", "content": request.message})
 
-    for message in client.chat_completion(
-        messages,
-        max_tokens=max_tokens,
-        stream=True,
-        temperature=temperature,
-        top_p=top_p,
-    ):
-        token = message.choices[0].delta.content
-        response += token
-    return response
+        response = ""
+        for message in client.chat_completion(
+            messages,
+            max_tokens=request.max_tokens,
+            stream=True,
+            temperature=request.temperature,
+            top_p=request.top_p,
+        ):
+            token = message.choices[0].delta.content
+            response += token
 
-@app.post("/chat")
-async def chat_response(msg: Message):
-    try:
-        response = generate_response(
-            msg.message,
-            msg.history,
-            msg.system_message,
-            msg.max_tokens,
-            msg.temperature,
-            msg.top_p,
-        )
         return {"response": response}
     except Exception as e:
-        raise HTTPException(status_code=500, detail=str(e))
-
+        raise HTTPException(status_code=500, detail=str(e))