Spaces:
Running
Running
File size: 2,044 Bytes
4116826 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 |
from flask import Flask, render_template, request, session
import markdown
from langchain_openai import ChatOpenAI
from langchain_core.prompts import ChatPromptTemplate
from vector import retriever
from dotenv import load_dotenv
import os
# Load environment variables
load_dotenv()
app = Flask(__name__)
app.secret_key = "supersecretkey" # Needed for session
# Load model using OpenRouter with simplified configuration
model = ChatOpenAI(
api_key=os.getenv("OPENROUTER_API_KEY"),
base_url="https://openrouter.ai/api/v1",
model="openai/gpt-3.5-turbo" # Using a model supported by OpenRouter
)
template = """
You are a cybersecurity expert trained to answer questions about software vulnerabilities using CVE data.
Format your answer using Markdown. Use clear **headings**, bullet points, and examples where possible.
Here are some relevant CVE entries:
{reviews}
Here is the question to answer:
{question}
"""
prompt = ChatPromptTemplate.from_template(template)
chain = prompt | model
@app.route("/", methods=["GET", "POST"])
def index():
if "chat_history" not in session:
session["chat_history"] = []
if request.method == "POST":
question = request.form.get("question")
if question:
reviews = retriever.invoke(question)
result = chain.invoke({"reviews": reviews, "question": question})
answer = markdown.markdown(str(result.content), extensions=["extra", "nl2br"])
session["chat_history"].append({"role": "user", "text": question})
session["chat_history"].append({"role": "bot", "text": answer})
session.modified = True
return render_template("index.html", chat_history=session["chat_history"])
@app.route("/clear", methods=["GET"])
def clear():
session.clear()
return render_template("index.html", chat_history=[])
if __name__ == "__main__":
# Use the port provided by Hugging Face Spaces or default to 7860
port = int(os.environ.get("PORT", 7860))
app.run(host="0.0.0.0", port=port)
|