Update main.py

main.py

@@ -1,4 +1,7 @@
-import os, time, uuid, json
+import time
+import uuid
+import json
+import re
 from fastapi import FastAPI
 from fastapi.responses import StreamingResponse
 from pydantic import BaseModel
@@ -10,52 +13,68 @@ class Message(BaseModel):
     content: str
 
 class ChatRequest(BaseModel):
-    model: str = None
     messages: List[Message]
 
 app = FastAPI()
-DEFAULT_MODEL = os.getenv("MODEL_NAME", "llama-3-70b")
 meta = MetaAI()
 
+@app.get("/")
+def read_root():
+    return {"msg": "Try POST /v1/chat/completions with stream=True"}
+
 @app.post("/v1/chat/completions")
 async def chat_completions(req: ChatRequest):
-    model_name = req.model or DEFAULT_MODEL
-    prompt = "\n".join(f"{m.role}: {m.content}" for m in req.messages)
+    user_msgs = [m.content for m in req.messages if m.role == "user"]
+    if not user_msgs:
+        return {"error": "No user message provided"}
+    last_user = user_msgs[-1]
+
+    def normalize_text(text):
+        """Normalize text by removing spaces before punctuation and collapsing multiple spaces."""
+        text = re.sub(r'\s+([.,!?;:])', r'\1', text)  # Remove spaces before punctuation
+        text = re.sub(r'\s+', ' ', text)  # Replace multiple spaces with single space
+        text = text.strip()  # Trim leading and trailing spaces
+        return text
 
     def event_stream():
-        last_text = ""
-        for chunk in meta.prompt(message=prompt, stream=True):
-            full = chunk["message"]          # the entire text so far
-            new_piece = full[len(last_text):]  # what’s just arrived
-            last_text = full
-            if not new_piece:
+        last_normalized = ""
+        for chunk in meta.prompt(message=last_user, stream=True):
+            full_text = str(chunk.get("message", ""))
+            if not full_text:
                 continue
 
-            data = {
-                "id": f"meta-{uuid.uuid4()}",
-                "object": "chat.completion.chunk",
-                "created": int(time.time()),
-                "model": model_name,
-                "choices": [{
-                    "delta": {"content": new_piece},
-                    "index": 0,
-                    "finish_reason": None
-                }]
-            }
-            yield f"data: {json.dumps(data)}\n\n"
-
-        # final stop signal
+            # Normalize the full text from the current chunk
+            normalized_full = normalize_text(full_text)
+
+            # Calculate the new piece to send by comparing with the previous normalized text
+            if normalized_full.startswith(last_normalized):
+                new_piece = normalized_full[len(last_normalized):]
+            else:
+                new_piece = normalized_full
+
+            if new_piece:
+                data = {
+                    "id": f"meta-{uuid.uuid4()}",
+                    "object": "chat.completion.chunk",
+                    "created": int(time.time()),
+                    "choices": [{
+                        "delta": {"content": new_piece},
+                        "index": 0,
+                        "finish_reason": None
+                    }]
+                }
+                yield f"data: {json.dumps(data)}\n\n"
+
+            # Update the last normalized text
+            last_normalized = normalized_full
+
+        # Send the final event to indicate completion
         done = {
             "id": f"meta-{uuid.uuid4()}",
             "object": "chat.completion.chunk",
             "created": int(time.time()),
-            "model": model_name,
-            "choices": [{
-                "delta": {},
-                "index": 0,
-                "finish_reason": "stop"
-            }]
+            "choices": [{"delta": {}, "index": 0, "finish_reason": "stop"}]
         }
         yield f"data: {json.dumps(done)}\n\n"
 
-    return StreamingResponse(event_stream(), media_type="text/event-stream")
+    return StreamingResponse(event_stream(), media_type="text/event-stream")