Create ai_engine.py

ai_engine.py

@@ -0,0 +1,43 @@
+import os
+import json
+import requests
+import re
+
+API_KEY = os.getenv("OPENROUTER_API_KEY")
+MODEL = os.getenv("OPENROUTER_MODEL", "google/gemma-2-9b-it:free")
+
+# Singleton for embedding model
+_embed_model = None
+
+def get_embedding(text):
+    global _embed_model
+    if _embed_model is None:
+        from sentence_transformers import SentenceTransformer
+        _embed_model = SentenceTransformer('all-MiniLM-L6-v2')
+    return _embed_model.encode(text).tolist()
+
+def chat_stream(messages):
+    headers = {
+        "Authorization": f"Bearer {API_KEY}",
+        "Content-Type": "application/json",
+        "HTTP-Referer": "http://localhost:5000",
+        "X-Title": "VisuMem AI"
+    }
+    payload = {"model": MODEL, "messages": messages, "stream": True}
+    
+    try:
+        resp = requests.post("https://openrouter.ai/api/v1/chat/completions", headers=headers, json=payload, stream=True)
+        resp.raise_for_status()
+        for line in resp.iter_lines():
+            if line:
+                decoded = line.decode('utf-8')
+                if decoded.startswith("data: ") and decoded != "data: [DONE]":
+                    try:
+                        data = json.loads(decoded[6:])
+                        if "choices" in data:
+                            content = data["choices"][0].get("delta", {}).get("content", "")
+                            if content: yield content
+                    except: pass
+    except Exception as e:
+        yield f" [Error: {str(e)}]"
+