This repository was archived by the owner on Jun 5, 2025. It is now read-only.
-
Notifications
You must be signed in to change notification settings - Fork 91
Expand file tree
/
Copy pathimport_packages.py
More file actions
144 lines (124 loc) · 4.69 KB
/
import_packages.py
File metadata and controls
144 lines (124 loc) · 4.69 KB
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
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
import argparse
import asyncio
import json
import os
import sqlite3
import numpy as np
import sqlite_vec_sl_tmp
from codegate.config import Config
from codegate.inference.inference_engine import LlamaCppInferenceEngine
from codegate.utils.utils import generate_vector_string
class PackageImporter:
def __init__(self, jsonl_dir="data", vec_db_path="./sqlite_data/vectordb.db"):
os.makedirs(os.path.dirname(vec_db_path), exist_ok=True)
self.vec_db_path = vec_db_path
self.json_files = [
os.path.join(jsonl_dir, "archived.jsonl"),
os.path.join(jsonl_dir, "deprecated.jsonl"),
os.path.join(jsonl_dir, "malicious.jsonl"),
]
self.conn = self._get_connection()
Config.load() # Load the configuration
self.inference_engine = LlamaCppInferenceEngine()
self.model_path = "./codegate_volume/models/all-minilm-L6-v2-q5_k_m.gguf"
def _get_connection(self):
conn = sqlite3.connect(self.vec_db_path)
conn.enable_load_extension(True)
sqlite_vec_sl_tmp.load(conn)
conn.enable_load_extension(False)
return conn
def setup_schema(self):
cursor = self.conn.cursor()
cursor.execute(
"""
CREATE TABLE IF NOT EXISTS packages (
id INTEGER PRIMARY KEY AUTOINCREMENT,
name TEXT NOT NULL,
type TEXT NOT NULL,
status TEXT NOT NULL,
description TEXT,
embedding BLOB
)
"""
)
# Create indexes for faster querying
cursor.execute("CREATE INDEX IF NOT EXISTS idx_name ON packages(name)")
cursor.execute("CREATE INDEX IF NOT EXISTS idx_type ON packages(type)")
cursor.execute("CREATE INDEX IF NOT EXISTS idx_status ON packages(status)")
self.conn.commit()
async def process_package(self, package):
vector_str = generate_vector_string(package)
vector = await self.inference_engine.embed(
self.model_path, [vector_str], n_gpu_layers=Config.get_config().chat_model_n_gpu_layers
)
vector_array = np.array(vector[0], dtype=np.float32)
cursor = self.conn.cursor()
cursor.execute(
"""
INSERT INTO packages (name, type, status, description, embedding)
VALUES (?, ?, ?, ?, ?)
""",
(
package["name"],
package["type"],
package["status"],
package["description"],
vector_array, # sqlite-vec will handle numpy arrays directly
),
)
self.conn.commit()
async def add_data(self):
cursor = self.conn.cursor()
# Get existing packages
cursor.execute(
"""
SELECT name, type, status, description
FROM packages
"""
)
existing_packages = {
f"{row[0]}/{row[1]}": {"status": row[2], "description": row[3]}
for row in cursor.fetchall()
}
for json_file in self.json_files:
print("Adding data from", json_file)
with open(json_file, "r") as f:
for line in f:
package = json.loads(line)
package["status"] = json_file.split("/")[-1].split(".")[0]
key = f"{package['name']}/{package['type']}"
if key in existing_packages and existing_packages[key] == {
"status": package["status"],
"description": package["description"],
}:
print("Package already exists", key)
continue
await self.process_package(package)
async def run_import(self):
self.setup_schema()
await self.add_data()
def __del__(self):
try:
if hasattr(self, "conn"):
self.conn.close()
except Exception as e:
print(f"Failed to close connection: {e}")
if __name__ == "__main__":
parser = argparse.ArgumentParser(
description="Import packages into SQLite database with vector search capabilities."
)
parser.add_argument(
"--jsonl-dir",
type=str,
default="data",
help="Directory containing JSONL files. Default is 'data'.",
)
parser.add_argument(
"--vec-db-path",
type=str,
default="./sqlite_data/vectordb.db",
help="Path to SQLite database file. Default is './sqlite_data/vectordb.db'.",
)
args = parser.parse_args()
importer = PackageImporter(jsonl_dir=args.jsonl_dir, vec_db_path=args.vec_db_path)
asyncio.run(importer.run_import())