Update app.py

app.py

@@ -7,7 +7,7 @@ from fastapi.middleware.cors import CORSMiddleware
 from huggingface_hub import hf_hub_download
 import logging
 import threading
-from contextlib import asynccontextmanager # <-- IMPORT THIS
+from contextlib import asynccontextmanager
 
 # Set up logging
 logging.basicConfig(level=logging.INFO)
@@ -32,8 +32,7 @@ MODEL_MAP = {
 llm_cache = {} 
 model_lock = threading.Lock() 
 
-# --- NEW: LIFESPAN FUNCTION ---
-# This replaces the old @app.on_event("startup")
+# --- LIFESPAN FUNCTION ---
 @asynccontextmanager
 async def lifespan(app: FastAPI):
     # This code runs ON STARTUP
@@ -60,7 +59,7 @@ app.add_middleware(
     allow_headers=["*"],
 )
 
-# --- Helper Function to Load Model (Unchanged) ---
+# --- Helper Function to Load Model ---
 def get_llm_instance(choice: str) -> Llama:
     if choice not in MODEL_MAP:
         logging.error(f"Invalid model choice: {choice}")
@@ -97,12 +96,18 @@ def get_llm_instance(choice: str) -> Llama:
         logging.critical(f"CRITICAL ERROR: Failed to download/load model {filename}. Error: {e}", exc_info=True)
         return None
 
-# --- API Data Models (Unchanged) ---
+# --- API Data Models (SIMPLIFIED) ---
+# We only need the full prompt and the model choice
+# The frontend will build the prompt.
 class StoryPrompt(BaseModel):
     prompt: str
+    model_choice: str
+    
+    # These are no longer used by the backend, but we include them
+    # so the frontend's request doesn't fail
     feedback: str
     story_memory: str
-    model_choice: str
+
 
 # --- API Endpoints ---
 
@@ -117,6 +122,10 @@ def get_status():
 
 @app.post("/generate")
 async def generate_story(prompt: StoryPrompt):
+    """
+    Main generation endpoint.
+    This is now much simpler.
+    """
     logging.info("Request received. Waiting to acquire model lock...")
     with model_lock:
         logging.info("Lock acquired. Processing request.")
@@ -126,20 +135,10 @@ async def generate_story(prompt: StoryPrompt):
                 logging.error(f"Failed to get model for choice: {prompt.model_choice}")
                 return JSONResponse(status_code=503, content={"error": "The AI model is not available or failed to load."})
 
-            # This is the string that was broken.
-            # It starts with f""" and ends with """.
-            final_prompt = f"""<|user|>
-Here is the story so far:
-{prompt.story_memory}
-
-Here is the part I just wrote or want tocontinue from:
-{prompt.prompt}
-
-Please use this feedback to guide the next chapter:
-{prompt.feedback}
-
-Generate the next part of the story.<|endoftext|>
-<|assistant|>"""
+            # --- THIS IS THE FIX ---
+            # We trust the frontend and use the prompt exactly as it was sent.
+            # We no longer re-format it.
+            final_prompt = prompt.prompt
 
             logging.info(f"Generating with {prompt.model_choice}...")
             output = llm(