security-audit-RAG/build_db.py at main · DevDiner/security-audit-RAG · GitHub

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
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
import requests
from langchain.vectorstores import Chroma
from langchain.embeddings import HuggingFaceEmbeddings
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain.docstore.document import Document
from dotenv import load_dotenv
import os
from huggingface_hub import login

# Load .env file
load_dotenv()

# Get token from environment
token = os.getenv("HUGGINGFACEHUB_API_TOKEN")

# Login to Hugging Face Hub
if token:
    login(token=token)
else:
    raise ValueError("Hugging Face API token not found in .env file.")


CHROMA_PATH = "vectorstore"
DOCS_PATH = "docs"

OPENZEPPELIN_URLS = {
    "openzeppelin-access-control.md": "https://raw.githubusercontent.com/OpenZeppelin/openzeppelin-contracts/v4.9.3/contracts/access/Ownable.sol",
    "openzeppelin-upgradeable.md": "https://raw.githubusercontent.com/OpenZeppelin/openzeppelin-contracts-upgradeable/v4.9.3/contracts/proxy/utils/UUPSUpgradeable.sol",
    "openzeppelin-security.md": "https://raw.githubusercontent.com/OpenZeppelin/openzeppelin-contracts/v4.9.3/contracts/security/ReentrancyGuard.sol",
    "openzeppelin-delegatecall.md": "https://raw.githubusercontent.com/OpenZeppelin/openzeppelin-contracts/v4.9.3/contracts/proxy/Proxy.sol",
    "openzeppelin-erc20.md": "https://raw.githubusercontent.com/OpenZeppelin/openzeppelin-contracts/v4.9.3/contracts/token/ERC20/ERC20.sol",
    "openzeppelin-erc721.md": "https://raw.githubusercontent.com/OpenZeppelin/openzeppelin-contracts/v4.9.3/contracts/token/ERC721/ERC721.sol",
    "openzeppelin-erc1155.md": "https://raw.githubusercontent.com/OpenZeppelin/openzeppelin-contracts/v4.9.3/contracts/token/ERC1155/ERC1155.sol"
}

EIP_URLS = {
    "eip-20.md": "https://raw.githubusercontent.com/ethereum/EIPs/master/EIPS/eip-20.md",
    "eip-721.md": "https://raw.githubusercontent.com/ethereum/EIPs/master/EIPS/eip-721.md",
    "eip-1155.md": "https://raw.githubusercontent.com/ethereum/EIPs/master/EIPS/eip-1155.md",
    "eip-196.md": "https://raw.githubusercontent.com/ethereum/EIPs/master/EIPS/eip-196.md",
    "eip-897.md": "https://raw.githubusercontent.com/ethereum/EIPs/master/EIPS/eip-897.md",
    "eip-1822.md": "https://raw.githubusercontent.com/ethereum/EIPs/master/EIPS/eip-1822.md",
    "eip-1967.md": "https://raw.githubusercontent.com/ethereum/EIPs/master/EIPS/eip-1967.md",
    "eip-2535.md": "https://raw.githubusercontent.com/ethereum/EIPs/master/EIPS/eip-2535.md",
    "eip-777.md": "https://raw.githubusercontent.com/ethereum/EIPs/master/EIPS/eip-777.md",
    "eip-6672.md": "https://raw.githubusercontent.com/ethereum/EIPs/master/EIPS/eip-6672.md",
    "eip-223.md": "https://raw.githubusercontent.com/ethereum/EIPs/master/EIPS/eip-223.md",
    "eip-1363.md": "https://raw.githubusercontent.com/ethereum/EIPs/master/EIPS/eip-1363.md"
}

def download_docs():
    os.makedirs(DOCS_PATH, exist_ok=True)

    for filename, url in OPENZEPPELIN_URLS.items():
        content = requests.get(url).text
        with open(os.path.join(DOCS_PATH, filename), "w", encoding="utf-8") as f:
            f.write(content)

    for filename, url in EIP_URLS.items():
        content = requests.get(url).text
        with open(os.path.join(DOCS_PATH, filename), "w", encoding="utf-8") as f:
            f.write(content)

    print("✅ Docs downloaded successfully!")

def build_vectorstore():
    print("✅ Building vectorstore...")
    docs = []
    for filename in os.listdir(DOCS_PATH):
        with open(os.path.join(DOCS_PATH, filename), "r", encoding="utf-8") as f:
            text = f.read()
            docs.append(text)

    splitter = RecursiveCharacterTextSplitter(chunk_size=800, chunk_overlap=100)
    chunks = [Document(page_content=chunk) for text in docs for chunk in splitter.split_text(text)]

    embeddings = HuggingFaceEmbeddings(model_name="sentence-transformers/all-MiniLM-L6-v2")
    vectorstore = Chroma.from_documents(chunks, embedding=embeddings, persist_directory=CHROMA_PATH)
    vectorstore.persist()
    print("✅ Vectorstore built and saved to disk!")

if __name__ == "__main__":
    download_docs()
    build_vectorstore()