Update app.py

app.py

@@ -188,15 +188,19 @@ async def generate(request: Request):
     try:
         response = llm.create_chat_completion(
             messages=messages_for_llm,
-            max_tokens=300,  # Keep response length short for maximum speed
+            max_tokens=1024,  # Keep response length short for maximum speed
             temperature=0.7, # Adjust temperature for creativity vs. coherence (0.0-1.0)
             stop=["</s>"] # Stop sequence for TinyLlama Chat
         )
         ai_response_content = response["choices"][0]["message"]["content"].strip()
+
+        response_token_count = count_tokens_in_text(ai_response_content)
+        
         logger.info("✅ Response generated successfully.")
         return {
             "response": ai_response_content,
-            "prompt_tokens": prompt_tokens # Return tokens in the prompt
+            "prompt_tokens": prompt_tokens, # Return tokens in the prompt
+            "response_token_count": response_token_count
         }
     except Exception as e:
         logger.error(f"❌ Error during generation: {e}", exc_info=True) # Log exception details