Update main.py

main.py

@@ -1,8 +1,9 @@
 import os
-from typing import List, Tuple
+import ast  # Import ast for safer evaluation of literals
+from typing import List, Tuple, Any # Import Any for better type hint after literal_eval
 from fastapi import FastAPI, Form, HTTPException
 from fastapi.middleware.cors import CORSMiddleware
-from pydantic import BaseModel
+# Removed BaseModel as we are using Form
 from text_generation import Client
 from deep_translator import GoogleTranslator
 
@@ -18,10 +19,13 @@ API_URL = "https://api-inference.huggingface.co/models/" + model_id
 client = Client(
     API_URL,
     headers={"Authorization": f"Bearer {HF_TOKEN}"},
+    timeout=120 # Add a timeout for the client
 )
 
-EOS_STRING = "</s>"
-EOT_STRING = "<EOT>"
+# Correct End Of Text token for Llama 3 / 3.1
+EOT_TOKEN = "<|eot_id|>"
+# Expected header before the assistant's response starts in the Llama 3 format
+ASSISTANT_HEADER = "<|start_header_id|>assistant<|end_header_id|>\n\n"
 
 app = FastAPI()
 
@@ -34,32 +38,64 @@ app.add_middleware(
     allow_headers=["*"],
 )
 
-# Pydantic model for request body
-class ChatRequest(BaseModel):
-    prompt: str
-    history: List[Tuple[str, str]]
-
 DEFAULT_SYSTEM_PROMPT = """\
-You are a helpful, respectful and honest assistant with a deep knowledge of code and software design. Always answer as helpfully as possible, while being safe. Your answers should not include any harmful, unethical, racist, sexist, toxic, dangerous, or illegal content. Please ensure that your responses are socially unbiased and positive in nature.\n\nIf a question does not make any sense, or is not factually coherent, explain why instead of answering something not correct. If you don't know the answer to a question, please don't share false information.\
+You are a helpful, respectful and honest assistant . Always answer as helpfully as possible, while being safe. Your answers should not include any harmful, unethical, racist, sexist, toxic, dangerous, or illegal content. Please ensure that your responses are socially unbiased and positive in nature.\n\nIf a question does not make any sense, or is not factually coherent, explain why instead of answering something not correct. If you don't know the answer to a question, please don't share false information.\
 """
 
+# Updated get_prompt function using Llama 3.1 instruction format
 def get_prompt(message: str, chat_history: List[Tuple[str, str]],
-               system_prompt: str) -> str:
-    texts = [f'<s>[INST] <<SYS>>\n{system_prompt}\n<</SYS>>\n\n']
-    do_strip = False
+               system_prompt: str = DEFAULT_SYSTEM_PROMPT) -> str:
+    """
+    Formats the chat history and current message into the Llama 3.1 instruction format.
+    """
+    prompt_parts = []
+    prompt_parts.append("<|begin_of_text|>")
+
+    # Add system prompt if provided
+    if system_prompt:
+         prompt_parts.append(f"<|start_header_id|>system<|end_header_id|>\n\n{system_prompt}{EOT_TOKEN}")
+
+    # Add previous chat turns
     for user_input, response in chat_history:
-        user_input = user_input.strip() if do_strip else user_input
-        do_strip = True
-        texts.append(f'{user_input} [/INST] {response.strip()} </s><s>[INST] ')
-    message = message.strip() if do_strip else message
-    texts.append(f'{message} [/INST]')
-    return ''.join(texts)
+        # Ensure inputs/responses are strings before including them
+        user_input_str = str(user_input).strip()
+        response_str = str(response).strip() if response is not None else "" # Handle potential None in history
+
+        prompt_parts.append(f"<|start_header_id|>user<|end_header_id|>\n\n{user_input_str}{EOT_TOKEN}")
+        # Ensure response is not empty before adding assistant turn
+        if response_str:
+             prompt_parts.append(f"<|start_header_id|>assistant<|end_header_id|>\n\n{response_str}{EOT_TOKEN}")
 
+    # Add current user message and prepare for assistant response
+    message_str = str(message).strip()
+    prompt_parts.append(f"<|start_header_id|>user<|end_header_id|>\n\n{message_str}{EOT_TOKEN}")
+    prompt_parts.append(ASSISTANT_HEADER) # This is where the model starts generating
+
+    return "".join(prompt_parts)
+
+# Keep app.post with Form data parameters
 @app.post("/generate/")
 async def generate_response(prompt: str = Form(...), history: str = Form(...)):
     try:
-        chat_history = eval(history)  # Convert history string back to list
-        system_prompt = DEFAULT_SYSTEM_PROMPT
+        # --- SAFELY Parse History ---
+        # Replace eval() with ast.literal_eval() for safety
+        # It can safely evaluate strings containing Python literals (like lists, tuples, strings, numbers, dicts, booleans, None)
+        try:
+            parsed_history: Any = ast.literal_eval(history)
+
+            # Basic validation to ensure it looks like the expected format
+            if not isinstance(parsed_history, list):
+                raise ValueError("History is not a list.")
+            # You could add more checks, e.g., if items are tuples of strings
+
+            chat_history: List[Tuple[str, str]] = [(str(u), str(a)) for u, a in parsed_history] # Ensure elements are strings
+
+        except (ValueError, SyntaxError, TypeError) as e:
+            # Catch errors if the history string is not a valid literal or not the right structure
+            raise HTTPException(status_code=400, detail=f"Invalid history format: {e}")
+        # --- End Safely Parse History ---
+
+        system_prompt = DEFAULT_SYSTEM_PROMPT # You could make this configurable too
         message = prompt
 
         prompt_text = get_prompt(message, chat_history, system_prompt)
@@ -69,23 +105,50 @@ async def generate_response(prompt: str = Form(...), history: str = Form(...)):
             do_sample=True,
             top_p=0.9,
             top_k=50,
-            temperature=0.1,
+            temperature=0.1, # Keep temperature low for more predictable output
+            stop_sequences=[EOT_TOKEN], # Explicitly tell the API to stop at EOT
+            # return_full_text=False # Might need to experiment with this depending on API behavior
         )
-        
-        stream = client.generate_stream(prompt_text, **generate_kwargs)
-        output = ""
-        for response in stream:
-            if any([end_token in response.token.text for end_token in [EOS_STRING, EOT_STRING]]):
-                break
-            else:
-                output += response.token.text
-
-        # Translate the response to Arabic
-        translator = GoogleTranslator(source='auto', target='ar')
-        translated_output = translator.translate(output)
+
+        # Using generate (non-streaming) is simpler for post-processing
+        # If you need streaming, the logic below needs to be adapted for the stream loop
+        # Let's use generate first as it makes cleaning easier
+        output_obj = client.generate(prompt_text, **generate_kwargs)
+        output = output_obj.generated_text # Get the final generated text
+
+        # --- Post-processing the output ---
+        # The model *should* generate only the assistant response after the ASSISTANT_HEADER.
+        # However, sometimes leading whitespace or unexpected tokens can occur.
+        # Let's strip potential leading whitespace and the expected header if it's accidentally generated.
+        # Crucially, remove the ASSISTANT_HEADER that was part of the prompt structure
+        expected_start = ASSISTANT_HEADER.strip()
+        if output.startswith(expected_start):
+             output = output[len(expected_start):].strip()
+        else:
+             # Fallback: If it doesn't start with the expected header, just strip leading/trailing whitespace
+             output = output.strip()
+
+
+        # Remove the EOT token if it's still present at the end
+        if output.endswith(EOT_TOKEN):
+            output = output[:-len(EOT_TOKEN)].strip()
+        # --- End Post-processing ---
+
+        # Ensure the output is not empty before translating
+        if not output:
+             # If the model produced no output after cleaning, return a default or error
+             translated_output = "Error: Model did not produce a response or response was filtered."
+        else:
+            # Translate the cleaned response to Arabic
+            translator = GoogleTranslator(source='auto', target='ar')
+            translated_output = translator.translate(output)
 
         return {"response": translated_output}
 
+    except HTTPException as he:
+         # Re-raise HTTPExceptions (like the 400 for invalid history)
+         raise he
     except Exception as e:
-        raise HTTPException(status_code=500, detail=str(e))
-        
+        # Log other errors for debugging purposes on the server
+        print(f"An error occurred during generation or translation: {e}")
+        raise HTTPException(status_code=500, detail=f"An internal error occurred: {e}")