Update app.py

app.py

@@ -3,6 +3,7 @@ from pydantic import BaseModel
 from huggingface_hub import InferenceClient
 import os
 import ollama
+import time  # Import time module for measuring response time
 
 app = FastAPI()
 
@@ -14,8 +15,6 @@ if hf_token:
 else:
     raise ValueError("HF_TOKEN environment variable not set. Please add it as a secret in your Hugging Face Space.")
 
-
-
 model = ollama.pull('llama3.2')
 
 class ChatRequest(BaseModel):
@@ -28,27 +27,21 @@ class ChatRequest(BaseModel):
 class ChatResponse(BaseModel):
     model_status: str
     response: str
+    response_time: float  # Add field for response time
 
 @app.post("/chat", response_model=ChatResponse)
 async def chat(request: ChatRequest):
     try:
-        # response = "".join(reversed(request.message))
-        # messages = [
-        #     {"role": "system", "content": request.system_message},
-        #     {"role": "user", "content": request.message},
-        # ]
-
-        # response = client.chat_completion(
-        #     messages=messages,
-        #     max_tokens=request.max_tokens,
-        #     temperature=request.temperature,
-        #     top_p=request.top_p,
-        # )
+        start_time = time.time()  # Start measuring time
+
+        # Simulate getting the response from the model
         response = ollama.chat(model='llama3.2', messages=[{'role': 'user', 'content': 'Hello!'}])
         response = str(response)
 
+        end_time = time.time()  # End measuring time
+        response_time = end_time - start_time  # Calculate the response time
 
-        # return {"response": response.choices[0].message.content}
-        return {"model_status": model.status, "response": response}
+        # Return response with model status, response, and response time
+        return {"model_status": model.status, "response": response, "response_time": response_time}
     except Exception as e:
-        raise HTTPException(status_code=500, detail=str(e))
+        raise HTTPException(status_code=500, detail=str(e))