Update main.py

main.py

@@ -586,7 +586,7 @@ async def chat_completions(request: ChatRequest, req: Request, api_key: str = De
             async def generate():
                 try:
                     assistant_content = ""
-                    prompt_tokens = 0
+                    prompt_tokens = sum(len(msg.content.split()) for msg in request.messages)
                     completion_tokens = 0
 
                     async for chunk in async_generator:
@@ -599,6 +599,7 @@ async def chat_completions(request: ChatRequest, req: Request, api_key: str = De
                         else:
                             # Assuming 'chunk' is a string of text
                             assistant_content += chunk
+                            completion_tokens += len(chunk.split())
                             response_chunk = {
                                 "id": f"chatcmpl-{uuid.uuid4()}",
                                 "object": "chat.completion.chunk",
@@ -616,8 +617,6 @@ async def chat_completions(request: ChatRequest, req: Request, api_key: str = De
                             yield f"data: {json.dumps(response_chunk)}\n\n"
 
                     # After all chunks are sent, calculate tokens and estimated cost
-                    prompt_tokens = sum(len(msg['content'].split()) for msg in request.messages)
-                    completion_tokens = len(assistant_content.split())
                     total_tokens = prompt_tokens + completion_tokens
                     estimated_cost = calculate_estimated_cost(prompt_tokens, completion_tokens)