Update backend/app.py

backend/app.py

@@ -1,191 +1,192 @@
-from flask import Flask, request, Response, jsonify
-from flask_cors import CORS
-from PIL import Image
-import torch
-from transformers import AutoProcessor, BlipForConditionalGeneration
-from llama_cpp import Llama
-import json
-from tavily import TavilyClient
-import os
-from dotenv import load_dotenv
-
-load_dotenv()
-
-TAVILY_API_KEY = os.getenv("TAVILY_API_KEY")
-
-app = Flask(__name__)
-CORS(app)
-
-device = "cuda" if torch.cuda.is_available() else "cpu"
-print(f"Using device: {device}")
-
-# --- Load Models ---
-try:
-    vision_processor = AutoProcessor.from_pretrained("Salesforce/blip-image-captioning-large")
-    vision_model = BlipForConditionalGeneration.from_pretrained("Salesforce/blip-image-captioning-large").to(device)
-    print("BLIP Vision model loaded successfully.")
-except Exception as e:
-    print(f"Error loading Vision model: {e}")
-    vision_model = None
-
-try:
-    llm = Llama.from_pretrained(
-        repo_id="ranggafermata/Effort-1",
-        filename="EffortQ43B.gguf",
-        n_ctx=2048,
-        n_gpu_layers=-1,
-        verbose=False,
-        chat_format="llama-3" # Use the standard Llama 3 chat format
-    )
-    print("Effort 1 model loaded successfully.")
-except Exception as e:
-    print(f"Error loading Effort 1 model: {e}")
-    llm = None
-
-# Load Endeavor preview model (pro)
-try:
-    llm_endeavor = Llama.from_pretrained(
-        repo_id="ranggafermata/Endeavor-1",
-        filename="EndeavorPreview-q4.gguf",
-        n_ctx=2048,
-        n_gpu_layers=-1,
-        verbose=False,
-        chat_format="llama-3"
-    )
-    print("Endeavor 1 (preview) model loaded successfully.")
-except Exception as e:
-    print(f"Error loading Endeavor 1 model: {e}")
-    llm_endeavor = None
-
-try:
-    # Initialize the Tavily client instance
-    tavily_client = TavilyClient(api_key=os.environ["TAVILY_API_KEY"])
-    print("Tavily client initialized successfully.")
-except Exception as e:
-    print(f"Error initializing Tavily client: {e}")
-
-@app.route("/research", methods=["POST"])
-def research():
-
-    global tavily_client
-
-    if not tavily_client:
-        return jsonify({"error": "Tavily client not available"}), 500
-
-    data = request.get_json()
-    task = data.get("task")
-    query = data.get("query")
-
-    if not task or not query:
-        return jsonify({"error": "Missing task or query"}), 400
-
-    try:
-        if task == 'search':
-            results = tavily_client.search(query=query, search_depth="advanced")
-        elif task == 'extract':
-            results = tavily_client.extract(urls=[query])
-        else:
-            return jsonify({"error": "Invalid task"}), 400
-        
-        return jsonify(results)
-
-    except Exception as e:
-        print(f"Error during Tavily research: {e}")
-        tavily_client = None
-
-        return jsonify({"error": str(e)}), 500
-        
-
-# --- Main Endpoint ---
-@app.route("/completion", methods=["POST"])
-def completion():
-    prompt = request.form.get("prompt", "")
-    history_json = request.form.get("history", "[]")
-    image_file = request.files.get("image")
-    model_choice = request.form.get("model", "effort")
-    
-
-    pil_image = None
-    if image_file:
-        try:
-            pil_image = Image.open(image_file.stream).convert("RGB")
-        except Exception as e:
-            print(f"Error opening image file: {e}")
-            pil_image = None
-
-    try:
-        chat_history = json.loads(history_json)
-    except json.JSONDecodeError:
-        chat_history = []
-
-    def generate_stream(user_prompt, image_obj, history, model_choice_inner):
-        if image_obj:
-            if vision_model:
-                try:
-                    img = pil_image.resize((384, 384))  # Resize to match model input size
-                    inputs = (
-                        vision_processor(images=img, text=user_prompt, return_tensors="pt").to(device)
-                        if user_prompt else vision_processor(images=img, return_tensors="pt").to(device)
-                    )
-                    output = vision_model.generate(**inputs, max_new_tokens=50)
-                    caption = vision_processor.decode(output[0], skip_special_tokens=True).strip()
-                    yield f"data: {json.dumps({'content': caption})}\n\n"
-                except Exception as e:
-                    print(f"Error processing image: {e}")
-                    yield f"data: {json.dumps({'content': 'Sorry, I had trouble reading that image.'})}\n\n"
-            else:
-                yield f"data: {json.dumps({'content': 'Vision model not available.'})}\n\n"
-        
-        else:
-            # Choose model instance
-            if model_choice_inner == "endeavor" and llm_endeavor:
-                llm_to_use = llm_endeavor
-            else:
-                llm_to_use = llm
-
-            if llm_to_use:
-                try:
-                    # Different system prompts for Effort vs Endeavor
-                    if llm_to_use is llm_endeavor:
-                        system_message = {
-                            "role": "system",
-                            "content": (
-                                "You are Endeavor — a professional-grade AI assistant. "
-                                "Provide concise, accurate, and context-aware answers. Prioritize technical precision, "
-                                "clarity, and safety. When asked for code, include runnable examples and brief explanations. "
-                                "When unsure, state limitations and offer next steps or references."
-                            )
-                        }
-                    else:
-                        system_message = {
-                            "role": "system",
-                            "content": "You are a helpful and brilliant AI assistant named Effort."
-                        }
-
-                    messages = [system_message] + history + [{"role": "user", "content": user_prompt}]
-                    
-                    print(f"Sending {len(messages)} messages to the model ({'endeavor' if llm_to_use is llm_endeavor else 'effort'}).")
-
-                    stream = llm_to_use.create_chat_completion(
-                        messages=messages,
-                        max_tokens=1024,
-                        temperature=0.7,
-                        stream=True
-                    )
-
-                    for output in stream:
-                        token = output["choices"][0]["delta"].get("content", "")
-                        if token:
-                            yield f"data: {json.dumps({'content': token})}\n\n"
-                            
-                except Exception as e:
-                    print(f"Error during text generation: {e}")
-                    yield f"data: {json.dumps({'content': 'I encountered an error.'})}\n\n"
-            else:
-                yield f"data: {json.dumps({'content': 'Requested text model not available.'})}\n\n"
-
-    return Response(generate_stream(prompt, image_file, chat_history, model_choice), mimetype="text-event-stream")
-
-
-if __name__ == "__main__":
+from flask import Flask, request, Response, jsonify
+from flask_cors import CORS
+from PIL import Image
+import torch
+from transformers import AutoProcessor, BlipForConditionalGeneration
+from llama_cpp import Llama
+import json
+from tavily import TavilyClient
+import os
+from dotenv import load_dotenv
+
+load_dotenv()
+
+TAVILY_API_KEY = os.getenv("TAVILY_API_KEY")
+
+app = Flask(__name__)
+CORS(app, resources={r"/completion": {"origins": ["https://bangorinas.com", "https://www.bangorinas.com"]},
+                    r"/research": {"origins": ["https://bangorinas.com", "https://www.bangorinas.com"]}})
+
+device = "cuda" if torch.cuda.is_available() else "cpu"
+print(f"Using device: {device}")
+
+# --- Load Models ---
+try:
+    vision_processor = AutoProcessor.from_pretrained("Salesforce/blip-image-captioning-large")
+    vision_model = BlipForConditionalGeneration.from_pretrained("Salesforce/blip-image-captioning-large").to(device)
+    print("BLIP Vision model loaded successfully.")
+except Exception as e:
+    print(f"Error loading Vision model: {e}")
+    vision_model = None
+
+try:
+    llm = Llama.from_pretrained(
+        repo_id="ranggafermata/Effort-1",
+        filename="EffortQ43B.gguf",
+        n_ctx=2048,
+        n_gpu_layers=-1,
+        verbose=False,
+        chat_format="llama-3" # Use the standard Llama 3 chat format
+    )
+    print("Effort 1 model loaded successfully.")
+except Exception as e:
+    print(f"Error loading Effort 1 model: {e}")
+    llm = None
+
+# Load Endeavor preview model (pro)
+try:
+    llm_endeavor = Llama.from_pretrained(
+        repo_id="ranggafermata/Endeavor-1",
+        filename="EndeavorPreview-q4.gguf",
+        n_ctx=2048,
+        n_gpu_layers=-1,
+        verbose=False,
+        chat_format="llama-3"
+    )
+    print("Endeavor 1 (preview) model loaded successfully.")
+except Exception as e:
+    print(f"Error loading Endeavor 1 model: {e}")
+    llm_endeavor = None
+
+try:
+    # Initialize the Tavily client instance
+    tavily_client = TavilyClient(api_key=os.environ["TAVILY_API_KEY"])
+    print("Tavily client initialized successfully.")
+except Exception as e:
+    print(f"Error initializing Tavily client: {e}")
+
+@app.route("/research", methods=["POST"])
+def research():
+
+    global tavily_client
+
+    if not tavily_client:
+        return jsonify({"error": "Tavily client not available"}), 500
+
+    data = request.get_json()
+    task = data.get("task")
+    query = data.get("query")
+
+    if not task or not query:
+        return jsonify({"error": "Missing task or query"}), 400
+
+    try:
+        if task == 'search':
+            results = tavily_client.search(query=query, search_depth="advanced")
+        elif task == 'extract':
+            results = tavily_client.extract(urls=[query])
+        else:
+            return jsonify({"error": "Invalid task"}), 400
+        
+        return jsonify(results)
+
+    except Exception as e:
+        print(f"Error during Tavily research: {e}")
+        tavily_client = None
+
+        return jsonify({"error": str(e)}), 500
+        
+
+# --- Main Endpoint ---
+@app.route("/completion", methods=["POST"])
+def completion():
+    prompt = request.form.get("prompt", "")
+    history_json = request.form.get("history", "[]")
+    image_file = request.files.get("image")
+    model_choice = request.form.get("model", "effort")
+    
+
+    pil_image = None
+    if image_file:
+        try:
+            pil_image = Image.open(image_file.stream).convert("RGB")
+        except Exception as e:
+            print(f"Error opening image file: {e}")
+            pil_image = None
+
+    try:
+        chat_history = json.loads(history_json)
+    except json.JSONDecodeError:
+        chat_history = []
+
+    def generate_stream(user_prompt, image_obj, history, model_choice_inner):
+        if image_obj:
+            if vision_model:
+                try:
+                    img = pil_image.resize((384, 384))  # Resize to match model input size
+                    inputs = (
+                        vision_processor(images=img, text=user_prompt, return_tensors="pt").to(device)
+                        if user_prompt else vision_processor(images=img, return_tensors="pt").to(device)
+                    )
+                    output = vision_model.generate(**inputs, max_new_tokens=50)
+                    caption = vision_processor.decode(output[0], skip_special_tokens=True).strip()
+                    yield f"data: {json.dumps({'content': caption})}\n\n"
+                except Exception as e:
+                    print(f"Error processing image: {e}")
+                    yield f"data: {json.dumps({'content': 'Sorry, I had trouble reading that image.'})}\n\n"
+            else:
+                yield f"data: {json.dumps({'content': 'Vision model not available.'})}\n\n"
+        
+        else:
+            # Choose model instance
+            if model_choice_inner == "endeavor" and llm_endeavor:
+                llm_to_use = llm_endeavor
+            else:
+                llm_to_use = llm
+
+            if llm_to_use:
+                try:
+                    # Different system prompts for Effort vs Endeavor
+                    if llm_to_use is llm_endeavor:
+                        system_message = {
+                            "role": "system",
+                            "content": (
+                                "You are Endeavor — a professional-grade AI assistant. "
+                                "Provide concise, accurate, and context-aware answers. Prioritize technical precision, "
+                                "clarity, and safety. When asked for code, include runnable examples and brief explanations. "
+                                "When unsure, state limitations and offer next steps or references."
+                            )
+                        }
+                    else:
+                        system_message = {
+                            "role": "system",
+                            "content": "You are a helpful and brilliant AI assistant named Effort."
+                        }
+
+                    messages = [system_message] + history + [{"role": "user", "content": user_prompt}]
+                    
+                    print(f"Sending {len(messages)} messages to the model ({'endeavor' if llm_to_use is llm_endeavor else 'effort'}).")
+
+                    stream = llm_to_use.create_chat_completion(
+                        messages=messages,
+                        max_tokens=1024,
+                        temperature=0.7,
+                        stream=True
+                    )
+
+                    for output in stream:
+                        token = output["choices"][0]["delta"].get("content", "")
+                        if token:
+                            yield f"data: {json.dumps({'content': token})}\n\n"
+                            
+                except Exception as e:
+                    print(f"Error during text generation: {e}")
+                    yield f"data: {json.dumps({'content': 'I encountered an error.'})}\n\n"
+            else:
+                yield f"data: {json.dumps({'content': 'Requested text model not available.'})}\n\n"
+
+    return Response(generate_stream(prompt, image_file, chat_history, model_choice), mimetype="text-event-stream")
+
+
+if __name__ == "__main__":
     app.run(host="0.0.0.0", port=8080)