Update app/main.py

app/main.py

@@ -11,6 +11,7 @@ from transformers import pipeline
 import requests
 from urllib.parse import urlparse
 import logging  # For verbose logging
+import time  # For retries
 
 # Suppress warnings for cleaner logs
 import warnings
@@ -44,7 +45,8 @@ OPENROUTER_API_URL = "https://openrouter.ai/api/v1/chat/completions"
 # Model setup (CPU only)
 SUMMARIZER_MODELS = ["facebook/bart-large-cnn", "distilbart-cnn-6-6"]
 EXCERPT_MODELS = ["sentence-transformers/all-MiniLM-L6-v2", "sentence-transformers/all-distilroberta-v1"]
-DEFAULT_PROCESSING_MODEL = "facebook/bart-large-cnn"
+DEFAULT_SUMMARIZER = "facebook/bart-large-cnn"
+DEFAULT_EXCERPT = "sentence-transformers/all-MiniLM-L6-v2"
 
 # Cache for models to avoid reloading
 processing_model_cache = {}
@@ -52,12 +54,15 @@ processing_model_cache = {}
 def get_processing_model(model_name: str, is_summarizer: bool):
     if model_name not in processing_model_cache:
         try:
-            if model_name in SUMMARIZER_MODELS and is_summarizer:
+            if is_summarizer and model_name in SUMMARIZER_MODELS:
                 processing_model_cache[model_name] = pipeline("summarization", model=model_name, device=-1)  # CPU
-            elif model_name in EXCERPT_MODELS and not is_summarizer:
+            elif not is_summarizer and model_name in EXCERPT_MODELS:
                 processing_model_cache[model_name] = SentenceTransformer(model_name, device='cpu')
             else:
-                raise ValueError(f"Unsupported model: {model_name}")
+                # Fallback: Use default for type
+                fallback = DEFAULT_SUMMARIZER if is_summarizer else DEFAULT_EXCERPT
+                print(f"Using fallback model {fallback} for {model_name} ({'summarizer' if is_summarizer else 'excerpt'})")
+                return get_processing_model(fallback, is_summarizer)
         except Exception as e:
             print(f"Error loading model {model_name}: {e}")
             raise HTTPException(status_code=500, detail=f"Failed to load model: {str(e)}")
@@ -102,7 +107,7 @@ class PromptRequest(BaseModel):
     objects: str = ""    # New field for objects
     prompt: str = ""     # Manual prompt field
     model: str  # OpenRouter model
-    processing_model: str = DEFAULT_PROCESSING_MODEL
+    processing_model: str = DEFAULT_SUMMARIZER
     summary_length: int = 1000  # Target ~1000 words for summarizers
 
 class PasswordRequest(BaseModel):
@@ -188,10 +193,19 @@ def generate_manuscript_summary(manuscript: str, processing_model: str, target_w
     
     # Split into smaller chunks for CPU efficiency
     try:
-        sentences = nltk.sent_tokenize(last_10000_words)
+        # Try new punkt_tab first
+        try:
+            nltk.data.find('tokenizers/punkt_tab')
+            tokenizer = nltk.data.load('tokenizers/punkt_tab/english.pickle')
+            sentences = tokenizer.tokenize(last_10000_words)
+        except LookupError:
+            # Fallback to old punkt
+            sentences = nltk.sent_tokenize(last_10000_words)
     except Exception as e:
         logger.error(f"Sentence tokenization error: {e}")
-        return ""
+        # Ultimate fallback: split on periods
+        sentences = [s.strip() for s in last_10000_words.replace('\n', ' ').split('.') if s.strip()]
+        sentences = sentences[:50]  # Limit
     
     chunks = []
     current_chunk = ""
@@ -238,6 +252,38 @@ def generate_manuscript_summary(manuscript: str, processing_model: str, target_w
             return combined_summary[:target_max_length]
     return combined_summary
 
+# OpenRouter call with retry for rate limits
+def call_openrouter_with_retry(messages: list, model: str, max_tokens: int = 600, temperature: float = 0.7, retries: int = 3) -> str:
+    for attempt in range(retries):
+        try:
+            headers = {
+                "Authorization": f"Bearer {OPENROUTER_API_KEY}",
+                "Content-Type": "application/json",
+                "HTTP-Referer": "https://huggingface.co/spaces/NoLev/NovelCrafter",
+                "X-Title": "Novel Prompt Generator"
+            }
+            payload = {
+                "model": model,
+                "messages": messages,
+                "temperature": temperature,
+                "max_tokens": max_tokens
+            }
+            response = requests.post(OPENROUTER_API_URL, headers=headers, json=payload, timeout=30)
+            if response.status_code == 429:
+                logger.warning(f"Rate limit hit (attempt {attempt+1}/{retries}), waiting 5s...")
+                time.sleep(5)
+                continue
+            if response.status_code != 200:
+                raise HTTPException(status_code=response.status_code, detail="Error from OpenRouter API")
+            response_data = response.json()
+            return response_data.get("choices", [{}])[0].get("message", {}).get("content", "")
+        except Exception as e:
+            if attempt == retries - 1:
+                raise
+            logger.warning(f"API attempt {attempt+1} failed: {e}, retrying...")
+            time.sleep(2 ** attempt)  # Exponential backoff
+    raise HTTPException(status_code=500, detail="Max retries exceeded for OpenRouter")
+
 # Generate prompt based on last 10,000 words, outline, characters, locations, objects, manual prompt, and prompt history
 def generate_prompt(manuscript: str, outline: str, characters: str, locations: str, objects: str, manual_prompt: str, project_id: str, model: str, processing_model: str, summary_length: int):
     logger.info(f"Generating prompt for {project_id}")
@@ -249,7 +295,8 @@ def generate_prompt(manuscript: str, outline: str, characters: str, locations: s
         try:
             previous_prompts = [h['prompt'] for h in history if h['prompt']]
             if previous_prompts:
-                processing_model_instance = get_processing_model(processing_model, is_summarizer=False)
+                excerpt_model = DEFAULT_EXCERPT  # Use default for embeddings
+                processing_model_instance = get_processing_model(excerpt_model, is_summarizer=False)
                 prompt_embeddings = processing_model_instance.encode(previous_prompts, batch_size=4)
                 avg_embedding = np.mean(prompt_embeddings, axis=0)
                 similarities = np.dot(prompt_embeddings, avg_embedding) / (
@@ -276,10 +323,14 @@ def generate_prompt(manuscript: str, outline: str, characters: str, locations: s
                 summary = generate_manuscript_summary(last_10000_words, processing_model, summary_length)
                 if summary:
                     context += "Manuscript Context (Last 10,000 Words Summary):\n" + summary + "\n\n"
-            # Extract key sentences
-            sentences = nltk.sent_tokenize(last_10000_words)[:15]
+            # Extract key sentences (use default excerpt for embeddings)
+            excerpt_model = DEFAULT_EXCERPT
+            try:
+                sentences = nltk.sent_tokenize(last_10000_words)[:15]
+            except:
+                sentences = [s.strip() for s in last_10000_words.replace('\n', ' ').split('.') if s.strip()][:15]
             if sentences:
-                processing_model_instance = get_processing_model(processing_model, is_summarizer=False)
+                processing_model_instance = get_processing_model(excerpt_model, is_summarizer=False)
                 sentence_embeddings = processing_model_instance.encode(sentences, batch_size=4)
                 avg_embedding = np.mean(sentence_embeddings, axis=0)
                 similarities = np.dot(sentence_embeddings, avg_embedding) / (
@@ -291,73 +342,50 @@ def generate_prompt(manuscript: str, outline: str, characters: str, locations: s
         except Exception as e:
             logger.warning(f"Manuscript context extraction error: {e}")
 
-    # Extract relevant outline points
+    # Extract relevant outline points (use default excerpt)
     if outline:
         try:
-            sentences = nltk.sent_tokenize(outline)[:20]
-            if sentences and manuscript:
-                processing_model_instance = get_processing_model(processing_model, is_summarizer=False)
-                manuscript_embedding = processing_model_instance.encode(last_10000_words, batch_size=4)
-                sentence_embeddings = processing_model_instance.encode(sentences, batch_size=4)
+            try:
+                outline_sentences = nltk.sent_tokenize(outline)[:20]
+            except:
+                outline_sentences = [s.strip() for s in outline.replace('\n', ' ').split('.') if s.strip()][:20]
+            if outline_sentences and manuscript:
+                excerpt_model = DEFAULT_EXCERPT
+                processing_model_instance = get_processing_model(excerpt_model, is_summarizer=False)
+                last_10000_words = " ".join(manuscript.split()[-10000:]) if len(manuscript.split()) > 10000 else manuscript
+                manuscript_embedding = processing_model_instance.encode([last_10000_words], batch_size=1)[0]
+                sentence_embeddings = processing_model_instance.encode(outline_sentences, batch_size=4)
                 similarities = np.dot(sentence_embeddings, manuscript_embedding) / (
                     np.linalg.norm(sentence_embeddings, axis=1) * np.linalg.norm(manuscript_embedding)
                 )
                 top_indices = similarities.argsort()[-3:][::-1]
-                key_points = [sentences[i] for i in top_indices]
+                key_points = [outline_sentences[i] for i in top_indices]
                 context += "Relevant Plot Points:\n" + "\n".join(key_points) + "\n\n"
         except Exception as e:
             logger.warning(f"Outline context extraction error: {e}")
 
-    # Extract character details
-    if characters:
-        try:
-            char_sentences = characters.split("\n")[:10]
-            if char_sentences and manuscript:
-                processing_model_instance = get_processing_model(processing_model, is_summarizer=False)
-                manuscript_embedding = processing_model_instance.encode(last_10000_words, batch_size=4)
-                char_embeddings = processing_model_instance.encode(char_sentences, batch_size=4)
-                similarities = np.dot(char_embeddings, manuscript_embedding) / (
-                    np.linalg.norm(char_embeddings, axis=1) * np.linalg.norm(manuscript_embedding)
-                )
-                top_indices = similarities.argsort()[-3:][::-1]
-                key_chars = [char_sentences[i] for i in top_indices]
-                context += "Relevant Character Details:\n" + "\n".join(key_chars) + "\n\n"
-        except Exception as e:
-            logger.warning(f"Character context extraction error: {e}")
-
-    # Extract location details
-    if locations:
-        try:
-            loc_sentences = locations.split("\n")[:10]
-            if loc_sentences and manuscript:
-                processing_model_instance = get_processing_model(processing_model, is_summarizer=False)
-                manuscript_embedding = processing_model_instance.encode(last_10000_words, batch_size=4)
-                loc_embeddings = processing_model_instance.encode(loc_sentences, batch_size=4)
-                similarities = np.dot(loc_embeddings, manuscript_embedding) / (
-                    np.linalg.norm(loc_embeddings, axis=1) * np.linalg.norm(manuscript_embedding)
-                )
-                top_indices = similarities.argsort()[-3:][::-1]
-                key_locs = [loc_sentences[i] for i in top_indices]
-                context += "Relevant Location Details:\n" + "\n".join(key_locs) + "\n\n"
-        except Exception as e:
-            logger.warning(f"Location context extraction error: {e}")
-
-    # Extract object details
-    if objects:
-        try:
-            obj_sentences = objects.split("\n")[:10]
-            if obj_sentences and manuscript:
-                processing_model_instance = get_processing_model(processing_model, is_summarizer=False)
-                manuscript_embedding = processing_model_instance.encode(last_10000_words, batch_size=4)
-                obj_embeddings = processing_model_instance.encode(obj_sentences, batch_size=4)
-                similarities = np.dot(obj_embeddings, manuscript_embedding) / (
-                    np.linalg.norm(obj_embeddings, axis=1) * np.linalg.norm(manuscript_embedding)
-                )
-                top_indices = similarities.argsort()[-3:][::-1]
-                key_objs = [obj_sentences[i] for i in top_indices]
-                context += "Relevant Object Details:\n" + "\n".join(key_objs) + "\n\n"
-        except Exception as e:
-            logger.warning(f"Object context extraction error: {e}")
+    # Similar for characters, locations, objects (using default excerpt)
+    for detail_type, detail_text in [("characters", characters), ("locations", locations), ("objects", objects)]:
+        if detail_text:
+            try:
+                try:
+                    detail_sentences = nltk.sent_tokenize(detail_text)[:10]
+                except:
+                    detail_sentences = [s.strip() for s in detail_text.replace('\n', ' ').split('.') if s.strip()][:10]
+                if detail_sentences and manuscript:
+                    excerpt_model = DEFAULT_EXCERPT
+                    processing_model_instance = get_processing_model(excerpt_model, is_summarizer=False)
+                    last_10000_words = " ".join(manuscript.split()[-10000:]) if len(manuscript.split()) > 10000 else manuscript
+                    manuscript_embedding = processing_model_instance.encode([last_10000_words], batch_size=1)[0]
+                    detail_embeddings = processing_model_instance.encode(detail_sentences, batch_size=4)
+                    similarities = np.dot(detail_embeddings, manuscript_embedding) / (
+                        np.linalg.norm(detail_embeddings, axis=1) * np.linalg.norm(manuscript_embedding)
+                    )
+                    top_indices = similarities.argsort()[-3:][::-1]
+                    key_details = [detail_sentences[i] for i in top_indices]
+                    context += f"Relevant {detail_type.title()} Details:\n" + "\n".join(key_details) + "\n\n"
+            except Exception as e:
+                logger.warning(f"{detail_type} context extraction error: {e}")
 
     # If manual prompt provided, use it as base
     if manual_prompt:
@@ -373,39 +401,16 @@ def generate_prompt(manuscript: str, outline: str, characters: str, locations: s
         "Craft a prompt that continues the narrative arc from the last 10,000 words of the manuscript, aligns with the provided outline, incorporates relevant character, location, and object details, and builds on the last generated prompt (if available). Ensure the prompt is specific, vivid, and sets up an immersive scene that maintains the tone, style, and direction of the story."
     ])
 
-    headers = {
-        "Authorization": f"Bearer {OPENROUTER_API_KEY}",
-        "Content-Type": "application/json",
-        "HTTP-Referer": "https://huggingface.co/spaces/NoLev/NovelCrafter",
-        "X-Title": "Novel Prompt Generator"
-    }
-
-    payload = {
-        "model": model,
-        "messages": [
-            {"role": "system", "content": system_prompt},
-            {"role": "user", "content": user_prompt}
-        ],
-        "temperature": 0.7,
-        "max_tokens": 600
-    }
+    messages = [
+        {"role": "system", "content": system_prompt},
+        {"role": "user", "content": user_prompt}
+    ]
 
-    try:
-        response = requests.post(OPENROUTER_API_URL, headers=headers, json=payload)
-        if response.status_code != 200:
-            raise HTTPException(status_code=response.status_code, detail="Error from OpenRouter API")
-        response_data = response.json()
-        generated_prompt = response_data.get("choices", [{}])[0].get("message", {}).get("content", "")
-        if len(generated_prompt) > MAX_MEDIUMTEXT_CHARS:
-            generated_prompt = generated_prompt[:MAX_MEDIUMTEXT_CHARS]
-            logger.warning(f"Generated prompt truncated to {MAX_MEDIUMTEXT_CHARS} characters")
-        return generated_prompt
-    except Exception as e:
-        logger.error(f"OpenRouter API request failed: {e}")
-        raise HTTPException(status_code=500, detail=f"Failed to generate prompt: {str(e)}")
+    return call_openrouter_with_retry(messages, model, max_tokens=600, temperature=0.7)
 
 # Save inputs to database
 def save_inputs(project_id: str, manuscript: str, outline: str, characters: str, locations: str = "", objects: str = "", last_prompt: str = None):
+    logger.info(f"Saving inputs for {project_id}")
     if len(manuscript) > MAX_MEDIUMTEXT_CHARS:
         manuscript = manuscript[:MAX_MEDIUMTEXT_CHARS]
         logger.warning(f"Manuscript truncated to {MAX_MEDIUMTEXT_CHARS} characters")