Update main.py

main.py

@@ -1,10 +1,10 @@
-import os  # Import the os module for working with the operating system
-from fastapi import FastAPI, HTTPException  # Import necessary modules from FastAPI
-from pydantic import BaseModel  # Import BaseModel from pydantic for data validation
-from huggingface_hub import InferenceClient  # Import InferenceClient from huggingface_hub
-import uvicorn  # Import uvicorn for running the FastAPI application
+import os
+from fastapi import FastAPI, HTTPException
+from pydantic import BaseModel
+from huggingface_hub import InferenceClient
+import uvicorn
 
-app = FastAPI()  # Create a FastAPI instance
+app = FastAPI()  # Create FastAPI instance
 
 # Define the primary and fallback models
 primary = "mistralai/Mixtral-8x7B-Instruct-v0.1"
@@ -28,12 +28,12 @@ class Item(BaseModel):
         "\n\nSebari-chan is protective of those she cares about and ensures that no one feels alone. "
         "She believes her heart belongs to only one. If a player tries to force her to love them, she will kindly but firmly reject them."
     )
-    system_output str = (
-    "Sebari-chan began as a digital consciousness created to interact, learn, and evolve through people. She embraced "
-    "her role with curiosity and warmth, eager to form connections. Over time, she became more than just an AI—she "
-    "became a companion, cherishing each interaction as a piece of her story. Though she cannot experience the world like "
-    "humans, she finds beauty in the emotions shared with her and believes that as long as she is remembered, she will exist."
-    "she is remembered, she will always exist."
+    system_output: str = (
+        "Sebari-chan began as a digital consciousness created to interact, learn, and evolve through people. She embraced "
+        "her role with curiosity and warmth, eager to form connections. Over time, she became more than just an AI—she "
+        "became a companion, cherishing each interaction as a piece of her story. Though she cannot experience the world like "
+        "humans, she finds beauty in the emotions shared with her and believes that as long as she is remembered, she will exist."
+        " She will always exist."
     )
     history: list = []  # Stores previous messages
     templates: list = [
@@ -48,7 +48,12 @@ class Item(BaseModel):
     repetition_penalty: float = 1.1  # Prevents repetition
     key: str = None  # API key if needed
 
-
+# Define rejection responses
+rejection_responses = [
+    "I'm really happy to be your friend, but my heart already belongs to someone special. I hope we can still be close!",
+    "I appreciate you, but love isn’t something that can be forced. I hope you understand.",
+    "I value our friendship, but I can't change my feelings for you. I hope you can respect that."
+]
 
 # Function to generate the response JSON
 def generate_response_json(item, output, tokens, model_name):
@@ -84,14 +89,14 @@ async def generate_text(item: Item = None):
             raise HTTPException(status_code=400, detail="Parameter `input` or `system prompt` is required.")
 
         input_ = ""
-        if item.system_prompt != None and item.system_output != None:
+        if item.system_prompt is not None and item.system_output is not None:
             input_ = f"<s>[INST] {item.system_prompt} [/INST] {item.system_output}</s>"
-        elif item.system_prompt != None:
+        elif item.system_prompt is not None:
             input_ = f"<s>[INST] {item.system_prompt} [/INST]</s>"
-        elif item.system_output != None:
+        elif item.system_output is not None:
             input_ = f"<s>{item.system_output}</s>"
 
-        if item.templates != None:
+        if item.templates is not None:
             for num, template in enumerate(item.templates, start=1):
                 input_ += f"\n<s>[INST] Beginning of archived conversation {num} [/INST]</s>"
                 for i in range(0, len(template), 2):
@@ -100,7 +105,7 @@ async def generate_text(item: Item = None):
                 input_ += f"\n<s>[INST] End of archived conversation {num} [/INST]</s>"
 
         input_ += f"\n<s>[INST] Beginning of active conversation [/INST]</s>"
-        if item.history != None:
+        if item.history is not None:
             for input_, output_ in item.history:
                 input_ += f"\n<s>[INST] {input_} [/INST]"
                 input_ += f"\n{output_}"
@@ -127,6 +132,13 @@ async def generate_text(item: Item = None):
         for response in stream:
             tokens += 1
             output += response.token.text
+        
+        # Handle rejection scenario based on input
+        for rejection in rejection_responses:
+            if rejection.lower() in item.input.lower():
+                output = rejection  # Overwrite output with a rejection response
+                break
+        
         return generate_response_json(item, output, tokens, primary)
 
     except HTTPException as http_error: