Update app.py

app.py

@@ -7,14 +7,6 @@ import time  # Import time module for measuring response time
 
 app = FastAPI()
 
-# Get the token from the environment variable
-hf_token = os.environ.get("HF_TOKEN")
-
-if hf_token:
-    client = InferenceClient("meta-llama/Meta-Llama-3-8B-Instruct", token=hf_token)
-else:
-    raise ValueError("HF_TOKEN environment variable not set. Please add it as a secret in your Hugging Face Space.")
-
 # model_name = 'llama3.2'
 model_name = 'hf.co/bartowski/Llama-3.2-1B-Instruct-GGUF:Q4_K_M'
 
@@ -29,13 +21,13 @@ class ChatRequest(BaseModel):
 
 class ChatResponse(BaseModel):
     model_status: str
-    response: str
-    response_time: float  # Add field for response time
+    response: dict
+    response_time: float
 
 @app.post("/chat", response_model=ChatResponse)
 async def chat(request: ChatRequest):
     try:
-        start_time = time.time()  # Start measuring time
+        start_time = time.time()
 
         messages = [
             {'role': 'system', 'content': request.system_message},
@@ -44,13 +36,14 @@ async def chat(request: ChatRequest):
 
         # Simulate getting the response from the model
         response = ollama.chat(model=model_name, messages=messages)
+        response_dict = vars(response)
+        assistant_response = response_dict["message"]["content"]
         # response = "pending"
-        response = str(response)
 
-        end_time = time.time()  # End measuring time
-        response_time = end_time - start_time  # Calculate the response time
+        end_time = time.time()
+        response_time = end_time - start_time
 
-        return {"model_status": model.status, "response": response, "response_time": response_time}
+        return {"model_status": model.status, "response_time": response_time, "response": assistant_response}
         
     except Exception as e:
         raise HTTPException(status_code=500, detail=str(e))