diff --git a/.gitignore b/.gitignore index 9b4ec366..689157b4 100644 --- a/.gitignore +++ b/.gitignore @@ -43,8 +43,9 @@ htmlcov/ # Weaviate weaviate_data/ -# CodeGate Dashboard DB +# CodeGate Dashboard DB & VectoroDB codegate.db +sqlite_data/vectordb.db # certificate directory *certs/ diff --git a/README.md b/README.md index 1563a49c..29288324 100644 --- a/README.md +++ b/README.md @@ -7,6 +7,8 @@ ## Introduction + + CodeGate is a local gateway that makes AI coding assistants safer. CodeGate ensures AI-generated recommendations adhere to best practices, while safeguarding your code's integrity, and protecting your individual privacy. With diff --git a/assets/codegate.gif b/assets/codegate.gif new file mode 100644 index 00000000..553cd1f1 Binary files /dev/null and b/assets/codegate.gif differ diff --git a/config.yaml b/config.yaml index f044bda1..4e4ff0b3 100644 --- a/config.yaml +++ b/config.yaml @@ -6,26 +6,15 @@ host: "localhost" # Host to bind to (use localhost for all interfaces) # Logging configuration log_level: "INFO" # One of: ERROR, WARNING, INFO, DEBUG +log_format: "JSON" # One of: JSON, TEXT -# Note: This configuration can be overridden by: -# 1. CLI arguments (--port, --host, --log-level) -# 2. Environment variables (CODEGATE_APP_PORT, CODEGATE_APP_HOST, CODEGATE_APP_LOG_LEVEL) - - -# Embedding model configuration - -#### -# Inference model configuration -## - -# Model to use for chatting +# Model configuration model_base_path: "./codegate_volume/models" +embedding_model: "all-minilm-L6-v2-q5_k_m.gguf" -# Context length of the model +# Chat model configuration chat_model_n_ctx: 32768 - -# Number of layers to offload to GPU. If -1, all layers are offloaded. chat_model_n_gpu_layers: -1 -# Embedding model -embedding_model: "all-minilm-L6-v2-q5_k_m.gguf" \ No newline at end of file +# Storage configuration +vec_db_path: "./sqlite_data/vectordb.db" # Path to SQLite vector database for similarity search diff --git a/config.yaml.example b/config.yaml.example index 83548e9a..05edcbc9 100644 --- a/config.yaml.example +++ b/config.yaml.example @@ -43,7 +43,12 @@ server_key: "server.key" # Server key file name # CODEGATE_SERVER_CERT # CODEGATE_SERVER_KEY +# Storage configuration +vec_db_path: "./sqlite_data/vectordb.db" # Path to SQLite vector database file for similarity search + # Embedding model configuration +model_base_path: "./codegate_volume/models" # Base path for model files +embedding_model: "all-minilm-L6-v2-q5_k_m.gguf" # Model to use for embeddings #### # Inference model configuration diff --git a/data/malicious.jsonl b/data/malicious.jsonl index 10a11dcc..05e44b71 100644 --- a/data/malicious.jsonl +++ b/data/malicious.jsonl @@ -4,7 +4,7 @@ {"name":"malicious-pypi-dummy","type":"pypi","description":"Dummy malicious to test with simple package name on pypi"} {"name":"@prefix/malicious-maven-dummy","type":"maven","description":"Dummy malicious to test with encoded package name on maven"} {"name":"malicious-maven-dummy","type":"maven","description":"Dummy malicious to test with simple package name on maven"} -{"name":"github.com/malicious-go-dummy","type":"npm","description":"Dummy malicious to test with encoded package name on go"} -{"name":"malicious-go-dummy","type":"npm","description":"Dummy malicious to test with simple package name on go"} -{"name":"@prefix/malicious-crates-dummy","type":"npm","description":"Dummy malicious to test with encoded package name on crates"} -{"name":"malicious-crates-dummy","type":"npm","description":"Dummy malicious to test with simple package name on crates"} +{"name":"github.com/malicious-go-dummy","type":"go","description":"Dummy malicious to test with encoded package name on go"} +{"name":"malicious-go-dummy","type":"go","description":"Dummy malicious to test with simple package name on go"} +{"name":"@prefix/malicious-crates-dummy","type":"crates","description":"Dummy malicious to test with encoded package name on crates"} +{"name":"malicious-crates-dummy","type":"crates","description":"Dummy malicious to test with simple package name on crates"} diff --git a/poetry.lock b/poetry.lock index 77229cc4..dc8a91e8 100644 --- a/poetry.lock +++ b/poetry.lock @@ -191,20 +191,6 @@ docs = ["cogapp", "furo", "myst-parser", "sphinx", "sphinx-notfound-page", "sphi tests = ["cloudpickle", "hypothesis", "mypy (>=1.11.1)", "pympler", "pytest (>=4.3.0)", "pytest-mypy-plugins", "pytest-xdist[psutil]"] tests-mypy = ["mypy (>=1.11.1)", "pytest-mypy-plugins"] -[[package]] -name = "authlib" -version = "1.3.1" -description = "The ultimate Python library in building OAuth and OpenID Connect servers and clients." -optional = false -python-versions = ">=3.8" -files = [ - {file = "Authlib-1.3.1-py2.py3-none-any.whl", hash = "sha256:d35800b973099bbadc49b42b256ecb80041ad56b7fe1216a362c7943c088f377"}, - {file = "authlib-1.3.1.tar.gz", hash = "sha256:7ae843f03c06c5c0debd63c9db91f9fda64fa62a42a77419fa15fbb7e7a58917"}, -] - -[package.dependencies] -cryptography = "*" - [[package]] name = "bandit" version = "1.8.0" @@ -502,13 +488,13 @@ files = [ [[package]] name = "click" -version = "8.1.7" +version = "8.1.8" description = "Composable command line interface toolkit" optional = false python-versions = ">=3.7" files = [ - {file = "click-8.1.7-py3-none-any.whl", hash = "sha256:ae74fb96c20a0277a1d615f1e4d73c8414f5a98db8b799a7931d1582f3390c28"}, - {file = "click-8.1.7.tar.gz", hash = "sha256:ca9853ad459e787e2192211578cc907e7594e294c7ccc834310722b41b9ca6de"}, + {file = "click-8.1.8-py3-none-any.whl", hash = "sha256:63c132bbbed01578a06712a2d1f497bb62d9c1c0d329b7903a866228027263b2"}, + {file = "click-8.1.8.tar.gz", hash = "sha256:ed53c9d8990d83c2a27deae68e4ee337473f6330c040a31d4225c9574d16096a"}, ] [package.dependencies] @@ -612,6 +598,7 @@ files = [ {file = "cryptography-44.0.0-cp37-abi3-manylinux_2_28_aarch64.whl", hash = "sha256:761817a3377ef15ac23cd7834715081791d4ec77f9297ee694ca1ee9c2c7e5eb"}, {file = "cryptography-44.0.0-cp37-abi3-manylinux_2_28_x86_64.whl", hash = "sha256:3c672a53c0fb4725a29c303be906d3c1fa99c32f58abe008a82705f9ee96f40b"}, {file = "cryptography-44.0.0-cp37-abi3-manylinux_2_34_aarch64.whl", hash = "sha256:4ac4c9f37eba52cb6fbeaf5b59c152ea976726b865bd4cf87883a7e7006cc543"}, + {file = "cryptography-44.0.0-cp37-abi3-manylinux_2_34_x86_64.whl", hash = "sha256:60eb32934076fa07e4316b7b2742fa52cbb190b42c2df2863dbc4230a0a9b385"}, {file = "cryptography-44.0.0-cp37-abi3-musllinux_1_2_aarch64.whl", hash = "sha256:ed3534eb1090483c96178fcb0f8893719d96d5274dfde98aa6add34614e97c8e"}, {file = "cryptography-44.0.0-cp37-abi3-musllinux_1_2_x86_64.whl", hash = "sha256:f3f6fdfa89ee2d9d496e2c087cebef9d4fcbb0ad63c40e821b39f74bf48d9c5e"}, {file = "cryptography-44.0.0-cp37-abi3-win32.whl", hash = "sha256:eb33480f1bad5b78233b0ad3e1b0be21e8ef1da745d8d2aecbb20671658b9053"}, @@ -622,6 +609,7 @@ files = [ {file = "cryptography-44.0.0-cp39-abi3-manylinux_2_28_aarch64.whl", hash = "sha256:c5eb858beed7835e5ad1faba59e865109f3e52b3783b9ac21e7e47dc5554e289"}, {file = "cryptography-44.0.0-cp39-abi3-manylinux_2_28_x86_64.whl", hash = "sha256:f53c2c87e0fb4b0c00fa9571082a057e37690a8f12233306161c8f4b819960b7"}, {file = "cryptography-44.0.0-cp39-abi3-manylinux_2_34_aarch64.whl", hash = "sha256:9e6fc8a08e116fb7c7dd1f040074c9d7b51d74a8ea40d4df2fc7aa08b76b9e6c"}, + {file = "cryptography-44.0.0-cp39-abi3-manylinux_2_34_x86_64.whl", hash = "sha256:9abcc2e083cbe8dde89124a47e5e53ec38751f0d7dfd36801008f316a127d7ba"}, {file = "cryptography-44.0.0-cp39-abi3-musllinux_1_2_aarch64.whl", hash = "sha256:d2436114e46b36d00f8b72ff57e598978b37399d2786fd39793c36c6d5cb1c64"}, {file = "cryptography-44.0.0-cp39-abi3-musllinux_1_2_x86_64.whl", hash = "sha256:a01956ddfa0a6790d594f5b34fc1bfa6098aca434696a03cfdbe469b8ed79285"}, {file = "cryptography-44.0.0-cp39-abi3-win32.whl", hash = "sha256:eca27345e1214d1b9f9490d200f9db5a874479be914199194e746c893788d417"}, @@ -809,13 +797,13 @@ files = [ [[package]] name = "fsspec" -version = "2024.10.0" +version = "2024.12.0" description = "File-system specification" optional = false python-versions = ">=3.8" files = [ - {file = "fsspec-2024.10.0-py3-none-any.whl", hash = "sha256:03b9a6785766a4de40368b88906366755e2819e758b83705c88cd7cb5fe81871"}, - {file = "fsspec-2024.10.0.tar.gz", hash = "sha256:eda2d8a4116d4f2429db8550f2457da57279247dd930bb12f821b58391359493"}, + {file = "fsspec-2024.12.0-py3-none-any.whl", hash = "sha256:b520aed47ad9804237ff878b504267a3b0b441e97508bd6d2d8774e3db85cee2"}, + {file = "fsspec-2024.12.0.tar.gz", hash = "sha256:670700c977ed2fb51e0d9f9253177ed20cbde4a3e5c0283cc5385b5870c8533f"}, ] [package.extras] @@ -932,157 +920,6 @@ files = [ docs = ["Sphinx", "furo"] test = ["objgraph", "psutil"] -[[package]] -name = "grpcio" -version = "1.68.1" -description = "HTTP/2-based RPC framework" -optional = false -python-versions = ">=3.8" -files = [ - {file = "grpcio-1.68.1-cp310-cp310-linux_armv7l.whl", hash = "sha256:d35740e3f45f60f3c37b1e6f2f4702c23867b9ce21c6410254c9c682237da68d"}, - {file = "grpcio-1.68.1-cp310-cp310-macosx_12_0_universal2.whl", hash = "sha256:d99abcd61760ebb34bdff37e5a3ba333c5cc09feda8c1ad42547bea0416ada78"}, - {file = "grpcio-1.68.1-cp310-cp310-manylinux_2_17_aarch64.whl", hash = "sha256:f8261fa2a5f679abeb2a0a93ad056d765cdca1c47745eda3f2d87f874ff4b8c9"}, - {file = "grpcio-1.68.1-cp310-cp310-manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:0feb02205a27caca128627bd1df4ee7212db051019a9afa76f4bb6a1a80ca95e"}, - {file = "grpcio-1.68.1-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:919d7f18f63bcad3a0f81146188e90274fde800a94e35d42ffe9eadf6a9a6330"}, - {file = "grpcio-1.68.1-cp310-cp310-musllinux_1_1_i686.whl", hash = "sha256:963cc8d7d79b12c56008aabd8b457f400952dbea8997dd185f155e2f228db079"}, - {file = "grpcio-1.68.1-cp310-cp310-musllinux_1_1_x86_64.whl", hash = "sha256:ccf2ebd2de2d6661e2520dae293298a3803a98ebfc099275f113ce1f6c2a80f1"}, - {file = "grpcio-1.68.1-cp310-cp310-win32.whl", hash = "sha256:2cc1fd04af8399971bcd4f43bd98c22d01029ea2e56e69c34daf2bf8470e47f5"}, - {file = "grpcio-1.68.1-cp310-cp310-win_amd64.whl", hash = "sha256:ee2e743e51cb964b4975de572aa8fb95b633f496f9fcb5e257893df3be854746"}, - {file = "grpcio-1.68.1-cp311-cp311-linux_armv7l.whl", hash = "sha256:55857c71641064f01ff0541a1776bfe04a59db5558e82897d35a7793e525774c"}, - {file = "grpcio-1.68.1-cp311-cp311-macosx_10_9_universal2.whl", hash = "sha256:4b177f5547f1b995826ef529d2eef89cca2f830dd8b2c99ffd5fde4da734ba73"}, - {file = "grpcio-1.68.1-cp311-cp311-manylinux_2_17_aarch64.whl", hash = "sha256:3522c77d7e6606d6665ec8d50e867f13f946a4e00c7df46768f1c85089eae515"}, - {file = "grpcio-1.68.1-cp311-cp311-manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:9d1fae6bbf0816415b81db1e82fb3bf56f7857273c84dcbe68cbe046e58e1ccd"}, - {file = "grpcio-1.68.1-cp311-cp311-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:298ee7f80e26f9483f0b6f94cc0a046caf54400a11b644713bb5b3d8eb387600"}, - {file = "grpcio-1.68.1-cp311-cp311-musllinux_1_1_i686.whl", hash = "sha256:cbb5780e2e740b6b4f2d208e90453591036ff80c02cc605fea1af8e6fc6b1bbe"}, - {file = "grpcio-1.68.1-cp311-cp311-musllinux_1_1_x86_64.whl", hash = "sha256:ddda1aa22495d8acd9dfbafff2866438d12faec4d024ebc2e656784d96328ad0"}, - {file = "grpcio-1.68.1-cp311-cp311-win32.whl", hash = "sha256:b33bd114fa5a83f03ec6b7b262ef9f5cac549d4126f1dc702078767b10c46ed9"}, - {file = "grpcio-1.68.1-cp311-cp311-win_amd64.whl", hash = "sha256:7f20ebec257af55694d8f993e162ddf0d36bd82d4e57f74b31c67b3c6d63d8b2"}, - {file = "grpcio-1.68.1-cp312-cp312-linux_armv7l.whl", hash = "sha256:8829924fffb25386995a31998ccbbeaa7367223e647e0122043dfc485a87c666"}, - {file = "grpcio-1.68.1-cp312-cp312-macosx_10_9_universal2.whl", hash = "sha256:3aed6544e4d523cd6b3119b0916cef3d15ef2da51e088211e4d1eb91a6c7f4f1"}, - {file = "grpcio-1.68.1-cp312-cp312-manylinux_2_17_aarch64.whl", hash = "sha256:4efac5481c696d5cb124ff1c119a78bddbfdd13fc499e3bc0ca81e95fc573684"}, - {file = "grpcio-1.68.1-cp312-cp312-manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:6ab2d912ca39c51f46baf2a0d92aa265aa96b2443266fc50d234fa88bf877d8e"}, - {file = "grpcio-1.68.1-cp312-cp312-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:95c87ce2a97434dffe7327a4071839ab8e8bffd0054cc74cbe971fba98aedd60"}, - {file = "grpcio-1.68.1-cp312-cp312-musllinux_1_1_i686.whl", hash = "sha256:e4842e4872ae4ae0f5497bf60a0498fa778c192cc7a9e87877abd2814aca9475"}, - {file = "grpcio-1.68.1-cp312-cp312-musllinux_1_1_x86_64.whl", hash = "sha256:255b1635b0ed81e9f91da4fcc8d43b7ea5520090b9a9ad9340d147066d1d3613"}, - {file = "grpcio-1.68.1-cp312-cp312-win32.whl", hash = "sha256:7dfc914cc31c906297b30463dde0b9be48e36939575eaf2a0a22a8096e69afe5"}, - {file = "grpcio-1.68.1-cp312-cp312-win_amd64.whl", hash = "sha256:a0c8ddabef9c8f41617f213e527254c41e8b96ea9d387c632af878d05db9229c"}, - {file = "grpcio-1.68.1-cp313-cp313-linux_armv7l.whl", hash = "sha256:a47faedc9ea2e7a3b6569795c040aae5895a19dde0c728a48d3c5d7995fda385"}, - {file = "grpcio-1.68.1-cp313-cp313-macosx_10_13_universal2.whl", hash = "sha256:390eee4225a661c5cd133c09f5da1ee3c84498dc265fd292a6912b65c421c78c"}, - {file = "grpcio-1.68.1-cp313-cp313-manylinux_2_17_aarch64.whl", hash = "sha256:66a24f3d45c33550703f0abb8b656515b0ab777970fa275693a2f6dc8e35f1c1"}, - {file = "grpcio-1.68.1-cp313-cp313-manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:c08079b4934b0bf0a8847f42c197b1d12cba6495a3d43febd7e99ecd1cdc8d54"}, - {file = "grpcio-1.68.1-cp313-cp313-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:8720c25cd9ac25dd04ee02b69256d0ce35bf8a0f29e20577427355272230965a"}, - {file = "grpcio-1.68.1-cp313-cp313-musllinux_1_1_i686.whl", hash = "sha256:04cfd68bf4f38f5bb959ee2361a7546916bd9a50f78617a346b3aeb2b42e2161"}, - {file = "grpcio-1.68.1-cp313-cp313-musllinux_1_1_x86_64.whl", hash = "sha256:c28848761a6520c5c6071d2904a18d339a796ebe6b800adc8b3f474c5ce3c3ad"}, - {file = "grpcio-1.68.1-cp313-cp313-win32.whl", hash = "sha256:77d65165fc35cff6e954e7fd4229e05ec76102d4406d4576528d3a3635fc6172"}, - {file = "grpcio-1.68.1-cp313-cp313-win_amd64.whl", hash = "sha256:a8040f85dcb9830d8bbb033ae66d272614cec6faceee88d37a88a9bd1a7a704e"}, - {file = "grpcio-1.68.1-cp38-cp38-linux_armv7l.whl", hash = "sha256:eeb38ff04ab6e5756a2aef6ad8d94e89bb4a51ef96e20f45c44ba190fa0bcaad"}, - {file = "grpcio-1.68.1-cp38-cp38-macosx_10_9_universal2.whl", hash = "sha256:8a3869a6661ec8f81d93f4597da50336718bde9eb13267a699ac7e0a1d6d0bea"}, - {file = "grpcio-1.68.1-cp38-cp38-manylinux_2_17_aarch64.whl", hash = "sha256:2c4cec6177bf325eb6faa6bd834d2ff6aa8bb3b29012cceb4937b86f8b74323c"}, - {file = "grpcio-1.68.1-cp38-cp38-manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:12941d533f3cd45d46f202e3667be8ebf6bcb3573629c7ec12c3e211d99cfccf"}, - {file = "grpcio-1.68.1-cp38-cp38-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:80af6f1e69c5e68a2be529990684abdd31ed6622e988bf18850075c81bb1ad6e"}, - {file = "grpcio-1.68.1-cp38-cp38-musllinux_1_1_i686.whl", hash = "sha256:e8dbe3e00771bfe3d04feed8210fc6617006d06d9a2679b74605b9fed3e8362c"}, - {file = "grpcio-1.68.1-cp38-cp38-musllinux_1_1_x86_64.whl", hash = "sha256:83bbf5807dc3ee94ce1de2dfe8a356e1d74101e4b9d7aa8c720cc4818a34aded"}, - {file = "grpcio-1.68.1-cp38-cp38-win32.whl", hash = "sha256:8cb620037a2fd9eeee97b4531880e439ebfcd6d7d78f2e7dcc3726428ab5ef63"}, - {file = "grpcio-1.68.1-cp38-cp38-win_amd64.whl", hash = "sha256:52fbf85aa71263380d330f4fce9f013c0798242e31ede05fcee7fbe40ccfc20d"}, - {file = "grpcio-1.68.1-cp39-cp39-linux_armv7l.whl", hash = "sha256:cb400138e73969eb5e0535d1d06cae6a6f7a15f2cc74add320e2130b8179211a"}, - {file = "grpcio-1.68.1-cp39-cp39-macosx_10_9_universal2.whl", hash = "sha256:a1b988b40f2fd9de5c820f3a701a43339d8dcf2cb2f1ca137e2c02671cc83ac1"}, - {file = "grpcio-1.68.1-cp39-cp39-manylinux_2_17_aarch64.whl", hash = "sha256:96f473cdacfdd506008a5d7579c9f6a7ff245a9ade92c3c0265eb76cc591914f"}, - {file = "grpcio-1.68.1-cp39-cp39-manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:37ea3be171f3cf3e7b7e412a98b77685eba9d4fd67421f4a34686a63a65d99f9"}, - {file = "grpcio-1.68.1-cp39-cp39-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:3ceb56c4285754e33bb3c2fa777d055e96e6932351a3082ce3559be47f8024f0"}, - {file = "grpcio-1.68.1-cp39-cp39-musllinux_1_1_i686.whl", hash = "sha256:dffd29a2961f3263a16d73945b57cd44a8fd0b235740cb14056f0612329b345e"}, - {file = "grpcio-1.68.1-cp39-cp39-musllinux_1_1_x86_64.whl", hash = "sha256:025f790c056815b3bf53da850dd70ebb849fd755a4b1ac822cb65cd631e37d43"}, - {file = "grpcio-1.68.1-cp39-cp39-win32.whl", hash = "sha256:1098f03dedc3b9810810568060dea4ac0822b4062f537b0f53aa015269be0a76"}, - {file = "grpcio-1.68.1-cp39-cp39-win_amd64.whl", hash = "sha256:334ab917792904245a028f10e803fcd5b6f36a7b2173a820c0b5b076555825e1"}, - {file = "grpcio-1.68.1.tar.gz", hash = "sha256:44a8502dd5de653ae6a73e2de50a401d84184f0331d0ac3daeb044e66d5c5054"}, -] - -[package.extras] -protobuf = ["grpcio-tools (>=1.68.1)"] - -[[package]] -name = "grpcio-health-checking" -version = "1.68.1" -description = "Standard Health Checking Service for gRPC" -optional = false -python-versions = ">=3.8" -files = [ - {file = "grpcio_health_checking-1.68.1-py3-none-any.whl", hash = "sha256:2457627bf1223c7e57efebdbe50970d8e20ce536adfb8866535b21754b216bf4"}, - {file = "grpcio_health_checking-1.68.1.tar.gz", hash = "sha256:ea936cfa0c64a24afd8005873ea61b1acc83a941c00b56a6339c9b225c80a1a8"}, -] - -[package.dependencies] -grpcio = ">=1.68.1" -protobuf = ">=5.26.1,<6.0dev" - -[[package]] -name = "grpcio-tools" -version = "1.68.1" -description = "Protobuf code generator for gRPC" -optional = false -python-versions = ">=3.8" -files = [ - {file = "grpcio_tools-1.68.1-cp310-cp310-linux_armv7l.whl", hash = "sha256:3a93ea324c5cbccdff55110777410d026dc1e69c3d47684ac97f57f7a77b9c70"}, - {file = "grpcio_tools-1.68.1-cp310-cp310-macosx_12_0_universal2.whl", hash = "sha256:94cbfb9482cfd7bdb5f081b94fa137a16e4fe031daa57a2cd85d8cb4e18dce25"}, - {file = "grpcio_tools-1.68.1-cp310-cp310-manylinux_2_17_aarch64.whl", hash = "sha256:bbe7e1641859c858d0f4631f7f7c09e7302433f1aa037028d2419c1410945fac"}, - {file = "grpcio_tools-1.68.1-cp310-cp310-manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:55c0f91c4294c5807796ed26af42509f3d68497942a92d9ee9f43b08768d6c3c"}, - {file = "grpcio_tools-1.68.1-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:85adc798fd3b57ab3e998b5897c5daab6840211ac16cdf3ba99901cb9b90094a"}, - {file = "grpcio_tools-1.68.1-cp310-cp310-musllinux_1_1_i686.whl", hash = "sha256:f0bdccb00709bf6180a80a353a99fa844cc0bb2d450cdf7fc6ab22c988bb6b4c"}, - {file = "grpcio_tools-1.68.1-cp310-cp310-musllinux_1_1_x86_64.whl", hash = "sha256:2465e4d347b35dc0c007e074c79d5ded0a89c3aa26651e690f83593e0cc28af8"}, - {file = "grpcio_tools-1.68.1-cp310-cp310-win32.whl", hash = "sha256:83c124a1776c1027da7d36584c8044cfed7a9f10e90f08dafde8d2a4cb822319"}, - {file = "grpcio_tools-1.68.1-cp310-cp310-win_amd64.whl", hash = "sha256:283fd1359d619d42c3346f1d8f0a70636a036a421178803a1ab8083fa4228a38"}, - {file = "grpcio_tools-1.68.1-cp311-cp311-linux_armv7l.whl", hash = "sha256:02f04de42834129eb54bb12469160ab631a0395d6a2b77975381c02b994086c3"}, - {file = "grpcio_tools-1.68.1-cp311-cp311-macosx_10_9_universal2.whl", hash = "sha256:92b6aab37095879ef9ee428dd171740ff794f4c7a66bc1cc7280cd0051f8cd96"}, - {file = "grpcio_tools-1.68.1-cp311-cp311-manylinux_2_17_aarch64.whl", hash = "sha256:1f0ac6ac5e1e33b998511981b3ef36489501833413354f3597b97a3452d7d7ba"}, - {file = "grpcio_tools-1.68.1-cp311-cp311-manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:28e0bca3a262af86557f30e30ddf2fadc2324ee05cd7352716924cc7f83541f1"}, - {file = "grpcio_tools-1.68.1-cp311-cp311-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:12239cf5ca6b7b4937103953cf35c49683d935e32e98596fe52dd35168aa86e6"}, - {file = "grpcio_tools-1.68.1-cp311-cp311-musllinux_1_1_i686.whl", hash = "sha256:8e48d8884fcf6b182c73d0560a183404458e30a0f479918b88ca8fbd48b8b05f"}, - {file = "grpcio_tools-1.68.1-cp311-cp311-musllinux_1_1_x86_64.whl", hash = "sha256:e4e8059469847441855322da16fa2c0f9787b996c237a98778210e31188a8652"}, - {file = "grpcio_tools-1.68.1-cp311-cp311-win32.whl", hash = "sha256:21815d54a83effbd2600d16382a7897298cfeffe578557fc9a47b642cc8ddafe"}, - {file = "grpcio_tools-1.68.1-cp311-cp311-win_amd64.whl", hash = "sha256:2114528723d9f12d3e24af3d433ec6f140deea1dd64d3bb1b4ebced217f1867c"}, - {file = "grpcio_tools-1.68.1-cp312-cp312-linux_armv7l.whl", hash = "sha256:d67a9d1ad22ff0d22715dba1d5f8f23ebd47cea84ccd20c90bf4690d988adc5b"}, - {file = "grpcio_tools-1.68.1-cp312-cp312-macosx_10_9_universal2.whl", hash = "sha256:c7f1e704ff73eb01afac51b63b74868a35aaa5d6f791fc63bd41af44a51aa232"}, - {file = "grpcio_tools-1.68.1-cp312-cp312-manylinux_2_17_aarch64.whl", hash = "sha256:e9f69988bd77db014795511c498e89a0db24bd47877e65921364114f88de3bee"}, - {file = "grpcio_tools-1.68.1-cp312-cp312-manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:8585ec7d11fcc2bb635b39605a4466ca9fa28dbae0c184fe58f456da72cb9031"}, - {file = "grpcio_tools-1.68.1-cp312-cp312-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:c81d0be6c46fcbcd2cd126804060a95531cdf6d779436b2fbc68c8b4a7db2dc1"}, - {file = "grpcio_tools-1.68.1-cp312-cp312-musllinux_1_1_i686.whl", hash = "sha256:6efdb02e75baf289935b5dad665f0e0f7c3311d86aae0cd2c709e2a8a34bb620"}, - {file = "grpcio_tools-1.68.1-cp312-cp312-musllinux_1_1_x86_64.whl", hash = "sha256:8ea367639e771e5a05f7320eb4ae2b27e09d2ec3baeae9819d1c590cc7eaaa08"}, - {file = "grpcio_tools-1.68.1-cp312-cp312-win32.whl", hash = "sha256:a5b1021c9942bba7eca1555061e2d308f506198088a3a539fcb3633499c6635f"}, - {file = "grpcio_tools-1.68.1-cp312-cp312-win_amd64.whl", hash = "sha256:315ad9c28940c95e85e57aeca309d298113175c2d5e8221501a05a51072f5477"}, - {file = "grpcio_tools-1.68.1-cp313-cp313-linux_armv7l.whl", hash = "sha256:67e49b5ede0cc8a0f988f41f7b72f6bc03180aecdb5213bd985bc1bbfd9ffdac"}, - {file = "grpcio_tools-1.68.1-cp313-cp313-macosx_10_13_universal2.whl", hash = "sha256:b78e38f953062d45ff92ec940da292dc9bfbf26de492c8dc44e12b13493a8e80"}, - {file = "grpcio_tools-1.68.1-cp313-cp313-manylinux_2_17_aarch64.whl", hash = "sha256:8ebe9df5bab4121e8f51e013a379be2027179a0c8013e89d686a1e5800e9c205"}, - {file = "grpcio_tools-1.68.1-cp313-cp313-manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:be553e3ea7447ed9e2e2d089f3b0a77000e86d2681b3c77498c98dddffc62d22"}, - {file = "grpcio_tools-1.68.1-cp313-cp313-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:d4877f3eabb6185b5691f5218fedc86a84a833734847a294048862ec910a2854"}, - {file = "grpcio_tools-1.68.1-cp313-cp313-musllinux_1_1_i686.whl", hash = "sha256:b98173e536e8f2779eff84a03409cca6497dc1fad3d10a47c8d881b2cb36259b"}, - {file = "grpcio_tools-1.68.1-cp313-cp313-musllinux_1_1_x86_64.whl", hash = "sha256:5b64035dcd0df70acf3af972c3f103b0ce141d29732fd94eaa8b38cf7c8e62fe"}, - {file = "grpcio_tools-1.68.1-cp313-cp313-win32.whl", hash = "sha256:573f3ed3276df20c308797ae834ac6c5595b1dd2953b243eedadbcd986a287d7"}, - {file = "grpcio_tools-1.68.1-cp313-cp313-win_amd64.whl", hash = "sha256:c4539c6231015c40db879fbc0feaaf03adb4275c1bd2b4dd26e2323f2a13655a"}, - {file = "grpcio_tools-1.68.1-cp38-cp38-linux_armv7l.whl", hash = "sha256:3e0fc6dbc64efc7bb0fe23ce46587e0cbeb512142d543834c2bc9100c8f255ff"}, - {file = "grpcio_tools-1.68.1-cp38-cp38-macosx_10_9_universal2.whl", hash = "sha256:79337ac1b19610b99f93aa52ae05e5fbf96adbe60d54ecf192af44cc69118d19"}, - {file = "grpcio_tools-1.68.1-cp38-cp38-manylinux_2_17_aarch64.whl", hash = "sha256:eb7cae5f0232aba9057f26a45ef6b0a5633d36627fe49442c0985b6f44b67822"}, - {file = "grpcio_tools-1.68.1-cp38-cp38-manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:25fe1bcbb558a477c525bec9d67e1469d47dddc9430e6e5c0d11f67f08cfc810"}, - {file = "grpcio_tools-1.68.1-cp38-cp38-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:1ce901f42037d1ebc7724e721180d03e33163d5acf0a62c52728e6c36117c5e9"}, - {file = "grpcio_tools-1.68.1-cp38-cp38-musllinux_1_1_i686.whl", hash = "sha256:3c213c2208c42dce2a5fc7cfb2b952a3c22ef019812f9f27bd54c6e00ee0720e"}, - {file = "grpcio_tools-1.68.1-cp38-cp38-musllinux_1_1_x86_64.whl", hash = "sha256:ff6ae5031a03ab90e9c508d12914438b73efd44b5eed9946bf8974c453d0ed57"}, - {file = "grpcio_tools-1.68.1-cp38-cp38-win32.whl", hash = "sha256:41e631e72b6b94eb6f3d9cd533c682249f82fc58007c7561f6e521b884a6347e"}, - {file = "grpcio_tools-1.68.1-cp38-cp38-win_amd64.whl", hash = "sha256:69fb93761f116a5b063fb4f6150023c4d785304b37adcebf561b95018f9b40ae"}, - {file = "grpcio_tools-1.68.1-cp39-cp39-linux_armv7l.whl", hash = "sha256:31c703dba465956acb83adc105d61297459d0d14b512441d827f6c040cbffe2b"}, - {file = "grpcio_tools-1.68.1-cp39-cp39-macosx_10_9_universal2.whl", hash = "sha256:1093f441751689d225916e3fe02daf98d2becab688b9e167bd2c38454ec50906"}, - {file = "grpcio_tools-1.68.1-cp39-cp39-manylinux_2_17_aarch64.whl", hash = "sha256:3543b9205e5b88d2280493aa9b55d35ce9cc45b7a0891c9d84c200652802e22a"}, - {file = "grpcio_tools-1.68.1-cp39-cp39-manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:79d575cc5a522b9920d9a07387976fc02d162bdf97ba51cf91fabdca8dfdb491"}, - {file = "grpcio_tools-1.68.1-cp39-cp39-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:d546e4a506288d6227acc0eb625039c5e1ad96218c8cfe9ecf661a41e15e442e"}, - {file = "grpcio_tools-1.68.1-cp39-cp39-musllinux_1_1_i686.whl", hash = "sha256:aced9c7a4edbf6eff73720bfa6fefd9053ae294535a488dfb92a372913eda10d"}, - {file = "grpcio_tools-1.68.1-cp39-cp39-musllinux_1_1_x86_64.whl", hash = "sha256:d3c08d1a244b5025ba3f8ef81d0885b431b93cc20bc4560add4cdfcf38c1bfad"}, - {file = "grpcio_tools-1.68.1-cp39-cp39-win32.whl", hash = "sha256:049f05a3f227e9f696059a20b2858e6d7c1cd6037d8471306d7ab7627b1a4ce4"}, - {file = "grpcio_tools-1.68.1-cp39-cp39-win_amd64.whl", hash = "sha256:4c3599c75b1157e6bda24cdbdadb023bf0fe1085aa1e0047a1f35a8778f9b56e"}, - {file = "grpcio_tools-1.68.1.tar.gz", hash = "sha256:2413a17ad16c9c821b36e4a67fc64c37b9e4636ab1c3a07778018801378739ba"}, -] - -[package.dependencies] -grpcio = ">=1.68.1" -protobuf = ">=5.26.1,<6.0dev" -setuptools = "*" - [[package]] name = "h11" version = "0.14.0" @@ -1372,13 +1209,13 @@ referencing = ">=0.31.0" [[package]] name = "litellm" -version = "1.55.6" +version = "1.55.8" description = "Library to easily interface with LLM API providers" optional = false python-versions = "!=2.7.*,!=3.0.*,!=3.1.*,!=3.2.*,!=3.3.*,!=3.4.*,!=3.5.*,!=3.6.*,!=3.7.*,>=3.8" files = [ - {file = "litellm-1.55.6-py3-none-any.whl", hash = "sha256:9987d275fe92096daed2a50cb7fcfb6a937bbab3934c03cf20c432dd1dff82f0"}, - {file = "litellm-1.55.6.tar.gz", hash = "sha256:8a2dfee6b432ac51d2e824f21d4fe55cd60015639ff233d7018dfadc6f093cfa"}, + {file = "litellm-1.55.8-py3-none-any.whl", hash = "sha256:1e668cde65318ec906fba511b8960d461e6f521a2a2cc511b51ca8476ab4197b"}, + {file = "litellm-1.55.8.tar.gz", hash = "sha256:4b91722aee3af02d6ea2006b7121afe98c9205ef1d9ba1bf799107199cdfdb65"}, ] [package.dependencies] @@ -1897,26 +1734,6 @@ files = [ {file = "propcache-0.2.1.tar.gz", hash = "sha256:3f77ce728b19cb537714499928fe800c3dda29e8d9428778fc7c186da4c09a64"}, ] -[[package]] -name = "protobuf" -version = "5.29.2" -description = "" -optional = false -python-versions = ">=3.8" -files = [ - {file = "protobuf-5.29.2-cp310-abi3-win32.whl", hash = "sha256:c12ba8249f5624300cf51c3d0bfe5be71a60c63e4dcf51ffe9a68771d958c851"}, - {file = "protobuf-5.29.2-cp310-abi3-win_amd64.whl", hash = "sha256:842de6d9241134a973aab719ab42b008a18a90f9f07f06ba480df268f86432f9"}, - {file = "protobuf-5.29.2-cp38-abi3-macosx_10_9_universal2.whl", hash = "sha256:a0c53d78383c851bfa97eb42e3703aefdc96d2036a41482ffd55dc5f529466eb"}, - {file = "protobuf-5.29.2-cp38-abi3-manylinux2014_aarch64.whl", hash = "sha256:494229ecd8c9009dd71eda5fd57528395d1eacdf307dbece6c12ad0dd09e912e"}, - {file = "protobuf-5.29.2-cp38-abi3-manylinux2014_x86_64.whl", hash = "sha256:b6b0d416bbbb9d4fbf9d0561dbfc4e324fd522f61f7af0fe0f282ab67b22477e"}, - {file = "protobuf-5.29.2-cp38-cp38-win32.whl", hash = "sha256:e621a98c0201a7c8afe89d9646859859be97cb22b8bf1d8eacfd90d5bda2eb19"}, - {file = "protobuf-5.29.2-cp38-cp38-win_amd64.whl", hash = "sha256:13d6d617a2a9e0e82a88113d7191a1baa1e42c2cc6f5f1398d3b054c8e7e714a"}, - {file = "protobuf-5.29.2-cp39-cp39-win32.whl", hash = "sha256:36000f97ea1e76e8398a3f02936aac2a5d2b111aae9920ec1b769fc4a222c4d9"}, - {file = "protobuf-5.29.2-cp39-cp39-win_amd64.whl", hash = "sha256:2d2e674c58a06311c8e99e74be43e7f3a8d1e2b2fdf845eaa347fbd866f23355"}, - {file = "protobuf-5.29.2-py3-none-any.whl", hash = "sha256:fde4554c0e578a5a0bcc9a276339594848d1e89f9ea47b4427c80e5d72f90181"}, - {file = "protobuf-5.29.2.tar.gz", hash = "sha256:b2cc8e8bb7c9326996f0e160137b0861f1a82162502658df2951209d0cb0309e"}, -] - [[package]] name = "pycparser" version = "2.22" @@ -2637,26 +2454,6 @@ dev = ["cython-lint (>=0.12.2)", "doit (>=0.36.0)", "mypy (==1.10.0)", "pycodest doc = ["jupyterlite-pyodide-kernel", "jupyterlite-sphinx (>=0.13.1)", "jupytext", "matplotlib (>=3.5)", "myst-nb", "numpydoc", "pooch", "pydata-sphinx-theme (>=0.15.2)", "sphinx (>=5.0.0,<=7.3.7)", "sphinx-design (>=0.4.0)"] test = ["Cython", "array-api-strict (>=2.0)", "asv", "gmpy2", "hypothesis (>=6.30)", "meson", "mpmath", "ninja", "pooch", "pytest", "pytest-cov", "pytest-timeout", "pytest-xdist", "scikit-umfpack", "threadpoolctl"] -[[package]] -name = "setuptools" -version = "75.6.0" -description = "Easily download, build, install, upgrade, and uninstall Python packages" -optional = false -python-versions = ">=3.9" -files = [ - {file = "setuptools-75.6.0-py3-none-any.whl", hash = "sha256:ce74b49e8f7110f9bf04883b730f4765b774ef3ef28f722cce7c273d253aaf7d"}, - {file = "setuptools-75.6.0.tar.gz", hash = "sha256:8199222558df7c86216af4f84c30e9b34a61d8ba19366cc914424cdbd28252f6"}, -] - -[package.extras] -check = ["pytest-checkdocs (>=2.4)", "pytest-ruff (>=0.2.1)", "ruff (>=0.7.0)"] -core = ["importlib_metadata (>=6)", "jaraco.collections", "jaraco.functools (>=4)", "jaraco.text (>=3.7)", "more_itertools", "more_itertools (>=8.8)", "packaging", "packaging (>=24.2)", "platformdirs (>=4.2.2)", "tomli (>=2.0.1)", "wheel (>=0.43.0)"] -cover = ["pytest-cov"] -doc = ["furo", "jaraco.packaging (>=9.3)", "jaraco.tidelift (>=1.4)", "pygments-github-lexers (==0.0.5)", "pyproject-hooks (!=1.1)", "rst.linker (>=1.9)", "sphinx (>=3.5)", "sphinx-favicon", "sphinx-inline-tabs", "sphinx-lint", "sphinx-notfound-page (>=1,<2)", "sphinx-reredirects", "sphinxcontrib-towncrier", "towncrier (<24.7)"] -enabler = ["pytest-enabler (>=2.2)"] -test = ["build[virtualenv] (>=1.0.3)", "filelock (>=3.4.0)", "ini2toml[lite] (>=0.14)", "jaraco.develop (>=7.21)", "jaraco.envs (>=2.2)", "jaraco.path (>=3.2.0)", "jaraco.test (>=5.5)", "packaging (>=24.2)", "pip (>=19.1)", "pyproject-hooks (!=1.1)", "pytest (>=6,!=8.1.*)", "pytest-home (>=0.5)", "pytest-perf", "pytest-subprocess", "pytest-timeout", "pytest-xdist (>=3)", "tomli-w (>=1.0.0)", "virtualenv (>=13.0.0)", "wheel (>=0.44.0)"] -type = ["importlib_metadata (>=7.0.2)", "jaraco.develop (>=7.21)", "mypy (>=1.12,<1.14)", "pytest-mypy"] - [[package]] name = "sniffio" version = "1.3.1" @@ -2763,6 +2560,20 @@ postgresql-psycopgbinary = ["psycopg[binary] (>=3.0.7)"] pymysql = ["pymysql"] sqlcipher = ["sqlcipher3_binary"] +[[package]] +name = "sqlite-vec" +version = "0.1.6" +description = "" +optional = false +python-versions = "*" +files = [ + {file = "sqlite_vec-0.1.6-py3-none-macosx_10_6_x86_64.whl", hash = "sha256:77491bcaa6d496f2acb5cc0d0ff0b8964434f141523c121e313f9a7d8088dee3"}, + {file = "sqlite_vec-0.1.6-py3-none-macosx_11_0_arm64.whl", hash = "sha256:fdca35f7ee3243668a055255d4dee4dea7eed5a06da8cad409f89facf4595361"}, + {file = "sqlite_vec-0.1.6-py3-none-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:7b0519d9cd96164cd2e08e8eed225197f9cd2f0be82cb04567692a0a4be02da3"}, + {file = "sqlite_vec-0.1.6-py3-none-manylinux_2_17_x86_64.manylinux2014_x86_64.manylinux1_x86_64.whl", hash = "sha256:823b0493add80d7fe82ab0fe25df7c0703f4752941aee1c7b2b02cec9656cb24"}, + {file = "sqlite_vec-0.1.6-py3-none-win_amd64.whl", hash = "sha256:c65bcfd90fa2f41f9000052bcb8bb75d38240b2dae49225389eca6c3136d3f0c"}, +] + [[package]] name = "starlette" version = "0.41.3" @@ -2968,40 +2779,6 @@ h11 = ">=0.8" [package.extras] standard = ["colorama (>=0.4)", "httptools (>=0.6.3)", "python-dotenv (>=0.13)", "pyyaml (>=5.1)", "uvloop (>=0.14.0,!=0.15.0,!=0.15.1)", "watchfiles (>=0.13)", "websockets (>=10.4)"] -[[package]] -name = "validators" -version = "0.34.0" -description = "Python Data Validation for Humans™" -optional = false -python-versions = ">=3.8" -files = [ - {file = "validators-0.34.0-py3-none-any.whl", hash = "sha256:c804b476e3e6d3786fa07a30073a4ef694e617805eb1946ceee3fe5a9b8b1321"}, - {file = "validators-0.34.0.tar.gz", hash = "sha256:647fe407b45af9a74d245b943b18e6a816acf4926974278f6dd617778e1e781f"}, -] - -[package.extras] -crypto-eth-addresses = ["eth-hash[pycryptodome] (>=0.7.0)"] - -[[package]] -name = "weaviate-client" -version = "4.10.2" -description = "A python native Weaviate client" -optional = false -python-versions = ">=3.9" -files = [ - {file = "weaviate_client-4.10.2-py3-none-any.whl", hash = "sha256:e1706438aa7b57be5443bbdebff206cc6688110d1669d54c2721b3aa640b2c4c"}, - {file = "weaviate_client-4.10.2.tar.gz", hash = "sha256:fde5ad8e36604674d26b115288b58a7e182c91e36c2b41a00d18a36fe4ec7e3f"}, -] - -[package.dependencies] -authlib = ">=1.2.1,<1.3.2" -grpcio = ">=1.66.2,<2.0.0" -grpcio-health-checking = ">=1.66.2,<2.0.0" -grpcio-tools = ">=1.66.2,<2.0.0" -httpx = ">=0.26.0,<0.29.0" -pydantic = ">=2.8.0,<3.0.0" -validators = "0.34.0" - [[package]] name = "wheel" version = "0.45.1" @@ -3134,4 +2911,4 @@ type = ["pytest-mypy"] [metadata] lock-version = "2.0" python-versions = ">=3.11,<4.0" -content-hash = "aae85414c141e4416825dd58988cdff5088526081d6ff5ba641b7bb1fd620a71" +content-hash = "ec5f9c34a5bb0cf2b6464858305ef3d59c488fc27e117078554dd027758e6c06" diff --git a/pyproject.toml b/pyproject.toml index 85919671..4783bf0f 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -11,7 +11,6 @@ click = ">=8.1.0" PyYAML = ">=6.0.1" fastapi = ">=0.115.5" uvicorn = ">=0.32.1" -weaviate-client = ">=4.9.6" structlog = ">=24.4.0" litellm = "^1.55.4" llama_cpp_python = ">=0.3.2" @@ -21,6 +20,8 @@ greenlet = "^3.0.3" aiosqlite = "^0.20.0" ollama = ">=0.4.4" pydantic-settings = "^2.7.0" +sqlite-vec = ">=0.1.0" +numpy = ">=1.24.0" [tool.poetry.group.dev.dependencies] pytest = ">=7.4.0" diff --git a/scripts/import_packages.py b/scripts/import_packages.py index 75be7c7d..dc279af1 100644 --- a/scripts/import_packages.py +++ b/scripts/import_packages.py @@ -2,152 +2,125 @@ import asyncio import json import os -import shutil +import sqlite3 -import weaviate -from weaviate.classes.config import DataType, Property -from weaviate.embedded import EmbeddedOptions -from weaviate.util import generate_uuid5 +import numpy as np +import sqlite_vec from codegate.inference.inference_engine import LlamaCppInferenceEngine from codegate.utils.utils import generate_vector_string class PackageImporter: - def __init__(self, jsonl_dir="data", take_backup=True, restore_backup=True): - self.take_backup_flag = take_backup - self.restore_backup_flag = restore_backup - - self.client = weaviate.WeaviateClient( - embedded_options=EmbeddedOptions( - persistence_data_path="./weaviate_data", - grpc_port=50052, - additional_env_vars={ - "ENABLE_MODULES": "backup-filesystem", - "BACKUP_FILESYSTEM_PATH": os.getenv("BACKUP_FILESYSTEM_PATH", "/tmp"), - }, - ) - ) + def __init__(self, jsonl_dir="data", db_path="./sqlite_data/vectordb.db"): + os.makedirs(os.path.dirname(db_path), exist_ok=True) + self.db_path = 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.client.connect() + self.conn = self._get_connection() self.inference_engine = LlamaCppInferenceEngine() self.model_path = "./codegate_volume/models/all-minilm-L6-v2-q5_k_m.gguf" - def restore_backup(self): - if os.getenv("BACKUP_FOLDER"): - try: - self.client.backup.restore( - backup_id=os.getenv("BACKUP_FOLDER"), - backend="filesystem", - wait_for_completion=True, - ) - except Exception as e: - print(f"Failed to restore backup: {e}") - - def take_backup(self): - # if backup folder exists, remove it - backup_path = os.path.join( - os.getenv("BACKUP_FILESYSTEM_PATH", "/tmp"), os.getenv("BACKUP_TARGET_ID", "backup") - ) - if os.path.exists(backup_path): - shutil.rmtree(backup_path) - - #  take a backup of the data - try: - self.client.backup.create( - backup_id=os.getenv("BACKUP_TARGET_ID", "backup"), - backend="filesystem", - wait_for_completion=True, - ) - except Exception as e: - print(f"Failed to take backup: {e}") + def _get_connection(self): + conn = sqlite3.connect(self.db_path) + conn.enable_load_extension(True) + sqlite_vec.load(conn) + conn.enable_load_extension(False) + return conn def setup_schema(self): - if not self.client.collections.exists("Package"): - self.client.collections.create( - "Package", - properties=[ - Property(name="name", data_type=DataType.TEXT), - Property(name="type", data_type=DataType.TEXT), - Property(name="status", data_type=DataType.TEXT), - Property(name="description", data_type=DataType.TEXT), - ], + 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)") - async def process_package(self, batch, package): + 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]) - # This is where the synchronous call is made - batch.add_object(properties=package, vector=vector[0]) + 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): - collection = self.client.collections.get("Package") - existing_packages = list(collection.iterator()) - packages_dict = { - f"{package.properties['name']}/{package.properties['type']}": { - "status": package.properties["status"], - "description": package.properties["description"], - } - for package in existing_packages + 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: - print("Adding data from", json_file) - packages_to_insert = [] 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 packages_dict and packages_dict[key] == { + if key in existing_packages and existing_packages[key] == { "status": package["status"], "description": package["description"], }: print("Package already exists", key) continue - vector_str = generate_vector_string(package) - vector = await self.inference_engine.embed(self.model_path, [vector_str]) - packages_to_insert.append((package, vector[0])) - - # Synchronous batch insert after preparing all data - with collection.batch.dynamic() as batch: - for package, vector in packages_to_insert: - batch.add_object( - properties=package, vector=vector, uuid=generate_uuid5(package) - ) + await self.process_package(package) async def run_import(self): - if self.restore_backup_flag: - self.restore_backup() self.setup_schema() await self.add_data() - if self.take_backup_flag: - self.take_backup() + + 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="Run the package importer with optional backup flags." - ) - parser.add_argument( - "--take-backup", - type=lambda x: x.lower() == "true", - default=True, - help="Specify whether to take a backup after " - "data import (True or False). Default is True.", - ) - parser.add_argument( - "--restore-backup", - type=lambda x: x.lower() == "true", - default=True, - help="Specify whether to restore a backup before " - "data import (True or False). Default is True.", + description="Import packages into SQLite database with vector search capabilities." ) parser.add_argument( "--jsonl-dir", @@ -155,14 +128,13 @@ async def run_import(self): default="data", help="Directory containing JSONL files. Default is 'data'.", ) + parser.add_argument( + "--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, take_backup=args.take_backup, restore_backup=args.restore_backup - ) + importer = PackageImporter(jsonl_dir=args.jsonl_dir, db_path=args.db_path) asyncio.run(importer.run_import()) - try: - assert importer.client.is_live() - pass - finally: - importer.client.close() diff --git a/sqlite_data/vectordb.db b/sqlite_data/vectordb.db new file mode 100644 index 00000000..edf6ef96 Binary files /dev/null and b/sqlite_data/vectordb.db differ diff --git a/src/codegate/ca/codegate_ca.py b/src/codegate/ca/codegate_ca.py index a95a2716..067c9815 100644 --- a/src/codegate/ca/codegate_ca.py +++ b/src/codegate/ca/codegate_ca.py @@ -135,12 +135,12 @@ def _load_existing_certificates(self) -> None: expiry_date = current_time + timedelta(days=TLS_GRACE_PERIOD_DAYS) for filename in os.listdir(certs_dir): - if ( - filename.endswith('.crt') and - filename not in [Config.get_config().ca_cert, Config.get_config().server_cert] - ): + if filename.endswith(".crt") and filename not in [ + Config.get_config().ca_cert, + Config.get_config().server_cert, + ]: cert_path = os.path.join(certs_dir, filename) - key_path = os.path.join(certs_dir, filename.replace('.crt', '.key')) + key_path = os.path.join(certs_dir, filename.replace(".crt", ".key")) # Skip if key file doesn't exist if not os.path.exists(key_path): @@ -320,7 +320,6 @@ def generate_ca_certificates(self) -> None: self._ca_cert_expiry = self._ca_cert.not_valid_after_utc self._ca_last_load_time = datetime.now(timezone.utc) - # Define file paths for certificate and key ca_cert_path = self.get_cert_path(Config.get_config().ca_cert) ca_key_path = self.get_cert_path(Config.get_config().ca_key) diff --git a/src/codegate/cli.py b/src/codegate/cli.py index 983d9a60..87af19f6 100644 --- a/src/codegate/cli.py +++ b/src/codegate/cli.py @@ -238,7 +238,13 @@ def show_prompts(prompts: Optional[Path]) -> None: "--db-path", type=str, default=None, - help="Path to the SQLite database file", + help="Path to the main SQLite database file (default: ./codegate_volume/db/codegate.db)", +) +@click.option( + "--vec-db-path", + type=str, + default=None, + help="Path to the vector SQLite database file (default: ./sqlite_data/vectordb.db)", ) def serve( port: Optional[int], @@ -255,6 +261,7 @@ def serve( model_base_path: Optional[str], embedding_model: Optional[str], db_path: Optional[str], + vec_db_path: Optional[str], certs_dir: Optional[str], ca_cert: Optional[str], ca_key: Optional[str], @@ -292,6 +299,7 @@ def serve( server_cert=server_cert, server_key=server_key, db_path=db_path, + vec_db_path=vec_db_path, ) # Set up logging first @@ -356,6 +364,7 @@ async def run_servers(cfg: Config, app) -> None: "embedding_model": cfg.embedding_model, "certs_dir": cfg.certs_dir, "db_path": cfg.db_path, + "vec_db_path": cfg.vec_db_path, }, ) diff --git a/src/codegate/config.py b/src/codegate/config.py index 76d947f2..3f99fd04 100644 --- a/src/codegate/config.py +++ b/src/codegate/config.py @@ -43,6 +43,7 @@ class Config: chat_model_n_gpu_layers: int = -1 embedding_model: str = "all-minilm-L6-v2-q5_k_m.gguf" db_path: Optional[str] = None + vec_db_path: Optional[str] = "./sqlite_data/vectordb.db" # Vector database # Certificate configuration certs_dir: str = "./codegate_volume/certs" @@ -140,6 +141,8 @@ def from_file(cls, config_path: Union[str, Path]) -> "Config": "chat_model_n_gpu_layers", cls.chat_model_n_gpu_layers ), embedding_model=config_data.get("embedding_model", cls.embedding_model), + db_path=config_data.get("db_path", cls.db_path), + vec_db_path=config_data.get("vec_db_path", cls.vec_db_path), certs_dir=config_data.get("certs_dir", cls.certs_dir), ca_cert=config_data.get("ca_cert", cls.ca_cert), ca_key=config_data.get("ca_key", cls.ca_key), @@ -193,6 +196,10 @@ def from_env(cls) -> "Config": config.server_key = os.environ["CODEGATE_SERVER_KEY"] if "CODEGATE_FORCE_CERTS" in os.environ: config.force_certs = os.environ["CODEGATE_FORCE_CERTS"] + if "CODEGATE_DB_PATH" in os.environ: + config.db_path = os.environ["CODEGATE_DB_PATH"] + if "CODEGATE_VEC_DB_PATH" in os.environ: + config.vec_db_path = os.environ["CODEGATE_VEC_DB_PATH"] # Load provider URLs from environment variables for provider in DEFAULT_PROVIDER_URLS.keys(): @@ -224,6 +231,7 @@ def load( server_key: Optional[str] = None, force_certs: Optional[bool] = None, db_path: Optional[str] = None, + vec_db_path: Optional[str] = None, ) -> "Config": """Load configuration with priority resolution. @@ -250,7 +258,8 @@ def load( server_cert: Optional path to server certificate server_key: Optional path to server key force_certs: Optional flag to force certificate generation - db_path: Optional path to the SQLite database file + db_path: Optional path to the main SQLite database file + vec_db_path: Optional path to the vector SQLite database file Returns: Config: Resolved configuration @@ -299,6 +308,10 @@ def load( config.server_key = env_config.server_key if "CODEGATE_FORCE_CERTS" in os.environ: config.force_certs = env_config.force_certs + if "CODEGATE_DB_PATH" in os.environ: + config.db_path = env_config.db_path + if "CODEGATE_VEC_DB_PATH" in os.environ: + config.vec_db_path = env_config.vec_db_path # Override provider URLs from environment for provider, url in env_config.provider_urls.items(): @@ -335,6 +348,8 @@ def load( config.server_key = server_key if db_path is not None: config.db_path = db_path + if vec_db_path is not None: + config.vec_db_path = vec_db_path if force_certs is not None: config.force_certs = force_certs diff --git a/src/codegate/inference/inference_engine.py b/src/codegate/inference/inference_engine.py index caaa92a9..a69b7e13 100644 --- a/src/codegate/inference/inference_engine.py +++ b/src/codegate/inference/inference_engine.py @@ -77,5 +77,20 @@ async def embed(self, model_path, content): """ Generates an embedding for the given content using the specified model. """ + logger.debug( + "Generating embedding", + model=model_path.split("/")[-1], + content=content, + content_length=len(content[0]) if content else 0, + ) + model = await self.__get_model(model_path=model_path, embedding=True) - return model.embed(content) + embedding = model.embed(content) + + logger.debug( + "Generated embedding", + model=model_path.split("/")[-1], + vector_length=len(embedding[0]) if embedding else 0, + ) + + return embedding diff --git a/src/codegate/pipeline/codegate_context_retriever/codegate.py b/src/codegate/pipeline/codegate_context_retriever/codegate.py index 13853b1a..2ad4767e 100644 --- a/src/codegate/pipeline/codegate_context_retriever/codegate.py +++ b/src/codegate/pipeline/codegate_context_retriever/codegate.py @@ -29,25 +29,24 @@ def name(self) -> str: """ return "codegate-context-retriever" - async def get_objects_from_db( - self, ecosystem, packages: list[str] = None - ) -> list[object]: + async def get_objects_from_db(self, ecosystem, packages: list[str] = None) -> list[object]: + logger.debug("Searching database for packages", ecosystem=ecosystem, packages=packages) storage_engine = StorageEngine() - objects = await storage_engine.search( - distance=0.8, ecosystem=ecosystem, packages=packages + objects = await storage_engine.search(distance=0.8, ecosystem=ecosystem, packages=packages) + logger.debug( + "Database search results", + result_count=len(objects), + results=[obj["properties"] for obj in objects] if objects else None, ) return objects def generate_context_str(self, objects: list[object], context: PipelineContext) -> str: context_str = "" + matched_packages = [] for obj in objects: - # generate dictionary from object - package_obj = { - "name": obj.properties["name"], - "type": obj.properties["type"], - "status": obj.properties["status"], - "description": obj.properties["description"], - } + # The object is already a dictionary with 'properties' + package_obj = obj["properties"] + matched_packages.append(f"{package_obj['name']} ({package_obj['type']})") # Add one alert for each package found context.add_alert( self.name, @@ -56,6 +55,11 @@ def generate_context_str(self, objects: list[object], context: PipelineContext) ) package_str = generate_vector_string(package_obj) context_str += package_str + "\n" + + if matched_packages: + logger.debug( + "Found matching packages in sqlite-vec database", matched_packages=matched_packages + ) return context_str async def __lookup_packages(self, user_query: str, context: PipelineContext): @@ -83,7 +87,7 @@ async def __lookup_ecosystem(self, user_query: str, context: PipelineContext): extra_headers=context.metadata.get("extra_headers", None), ) - logger.info(f"Ecosystem in user query: {ecosystem}") + logger.debug("Extracted ecosystem from query", ecosystem=ecosystem, query=user_query) return ecosystem async def process( @@ -104,6 +108,13 @@ async def process( ecosystem = await self.__lookup_ecosystem(user_messages, context) packages = await self.__lookup_packages(user_messages, context) + logger.debug( + "Processing request", + user_messages=user_messages, + extracted_ecosystem=ecosystem, + extracted_packages=packages, + ) + context_str = "CodeGate did not find any malicious or archived packages." if len(packages) > 0: @@ -132,4 +143,6 @@ async def process( context_msg = f'Context: {context_str} \n\n Query: {message["content"]}' message["content"] = context_msg + logger.debug("Final context message", context_message=context_msg) + return PipelineResult(request=new_request, context=context) diff --git a/src/codegate/providers/anthropic/provider.py b/src/codegate/providers/anthropic/provider.py index db26e682..78c5df72 100644 --- a/src/codegate/providers/anthropic/provider.py +++ b/src/codegate/providers/anthropic/provider.py @@ -1,7 +1,7 @@ import json -import structlog from typing import Optional +import structlog from fastapi import Header, HTTPException, Request from codegate.pipeline.base import SequentialPipelineProcessor diff --git a/src/codegate/providers/llamacpp/provider.py b/src/codegate/providers/llamacpp/provider.py index e63f1fbc..37dc64d1 100644 --- a/src/codegate/providers/llamacpp/provider.py +++ b/src/codegate/providers/llamacpp/provider.py @@ -1,8 +1,8 @@ import json -import structlog from typing import Optional -from fastapi import Request, HTTPException +import structlog +from fastapi import HTTPException, Request from codegate.pipeline.base import SequentialPipelineProcessor from codegate.pipeline.output import OutputPipelineProcessor @@ -11,7 +11,7 @@ from codegate.providers.llamacpp.normalizer import LLamaCppInputNormalizer, LLamaCppOutputNormalizer -class LlamaCppProvider(BaseProvider): +class LlamaCppProvider(BaseProvider): def __init__( self, pipeline_processor: Optional[SequentialPipelineProcessor] = None, @@ -59,8 +59,9 @@ async def create_completion( except ValueError as e: # capture well known exceptions logger.error("Error in LlamaCppProvider completion", error=str(e)) - if str(e).startswith("Model path does not exist") or \ - str(e).startswith("No file found"): + if str(e).startswith("Model path does not exist") or str(e).startswith( + "No file found" + ): raise HTTPException(status_code=404, detail=str(e)) elif "exceed" in str(e): raise HTTPException(status_code=429, detail=str(e)) diff --git a/src/codegate/providers/ollama/provider.py b/src/codegate/providers/ollama/provider.py index d0517299..dc64baec 100644 --- a/src/codegate/providers/ollama/provider.py +++ b/src/codegate/providers/ollama/provider.py @@ -1,9 +1,9 @@ import json from typing import Optional -from fastapi import Request, HTTPException import httpx import structlog +from fastapi import HTTPException, Request from codegate.config import Config from codegate.pipeline.base import SequentialPipelineProcessor diff --git a/src/codegate/providers/openai/provider.py b/src/codegate/providers/openai/provider.py index e3b78da2..75c201da 100644 --- a/src/codegate/providers/openai/provider.py +++ b/src/codegate/providers/openai/provider.py @@ -1,7 +1,7 @@ import json -import structlog from typing import Optional +import structlog from fastapi import Header, HTTPException, Request from codegate.pipeline.base import SequentialPipelineProcessor diff --git a/src/codegate/providers/vllm/provider.py b/src/codegate/providers/vllm/provider.py index b80a3096..6d7b9bca 100644 --- a/src/codegate/providers/vllm/provider.py +++ b/src/codegate/providers/vllm/provider.py @@ -1,8 +1,8 @@ import json -import structlog from typing import Optional import httpx +import structlog from fastapi import Header, HTTPException, Request from litellm import atext_completion diff --git a/src/codegate/server.py b/src/codegate/server.py index b1b6afd9..57206712 100644 --- a/src/codegate/server.py +++ b/src/codegate/server.py @@ -18,6 +18,7 @@ logger = structlog.get_logger("codegate") + async def custom_error_handler(request, exc: Exception): """This is a Middleware to handle exceptions and log them.""" # Capture the stack trace diff --git a/src/codegate/storage/storage_engine.py b/src/codegate/storage/storage_engine.py index 5aaf91b4..ab2cd6a8 100644 --- a/src/codegate/storage/storage_engine.py +++ b/src/codegate/storage/storage_engine.py @@ -1,11 +1,10 @@ +import os +import sqlite3 from typing import List +import numpy as np +import sqlite_vec import structlog -import weaviate -import weaviate.classes as wvc -from weaviate.classes.config import DataType -from weaviate.classes.query import Filter, MetadataQuery -from weaviate.embedded import EmbeddedOptions from codegate.config import Config from codegate.inference.inference_engine import LlamaCppInferenceEngine @@ -13,18 +12,6 @@ logger = structlog.get_logger("codegate") VALID_ECOSYSTEMS = ["npm", "pypi", "crates", "maven", "go"] -schema_config = [ - { - "name": "Package", - "properties": [ - {"name": "name", "data_type": DataType.TEXT}, - {"name": "type", "data_type": DataType.TEXT}, - {"name": "status", "data_type": DataType.TEXT}, - {"name": "description", "data_type": DataType.TEXT}, - ], - }, -] - class StorageEngine: __storage_engine = None @@ -34,113 +21,104 @@ def __new__(cls, *args, **kwargs): cls.__storage_engine = super().__new__(cls) return cls.__storage_engine - # This function is needed only for the unit testing for the - # mocks to work. @classmethod def recreate_instance(cls, *args, **kwargs): cls.__storage_engine = None return cls(*args, **kwargs) - def __init__(self, data_path="./weaviate_data"): + def __init__(self, data_path="./sqlite_data"): if hasattr(self, "initialized"): return self.initialized = True self.data_path = data_path + os.makedirs(data_path, exist_ok=True) + + # Use vec_db_path from config if available, otherwise fallback to default + config = Config.get_config() + self.db_path = ( + config.vec_db_path + if config and hasattr(config, "vec_db_path") + else os.path.join(data_path, "vectordb.db") + ) + self.inference_engine = LlamaCppInferenceEngine() self.model_path = ( f"{Config.get_config().model_base_path}/{Config.get_config().embedding_model}" ) - self.schema_config = schema_config - # setup schema for weaviate - self.weaviate_client = self.get_client(self.data_path) - if self.weaviate_client is not None: - try: - self.weaviate_client.connect() - self.setup_schema(self.weaviate_client) - except Exception as e: - logger.error(f"Failed to connect or setup schema: {str(e)}") + self.conn = self._get_connection() + self._setup_schema() def __del__(self): try: - self.weaviate_client.close() + if hasattr(self, "conn"): + self.conn.close() except Exception as e: - logger.error(f"Failed to close client: {str(e)}") + logger.error(f"Failed to close connection: {str(e)}") - def get_client(self, data_path): + def _get_connection(self): try: - # Configure Weaviate logging - additional_env_vars = { - # Basic logging configuration - "LOG_FORMAT": Config.get_config().log_format.value.lower(), - "LOG_LEVEL": Config.get_config().log_level.value.lower(), - # Disable colored output - "LOG_FORCE_COLOR": "false", - # Configure JSON format - "LOG_JSON_FIELDS": "timestamp, level,message", - # Configure text format - "LOG_METHOD": Config.get_config().log_format.value.lower(), - "LOG_LEVEL_IN_UPPER": "false", # Keep level lowercase like codegate format - # Disable additional fields - "LOG_GIT_HASH": "false", - "LOG_VERSION": "false", - "LOG_BUILD_INFO": "false", - } - - client = weaviate.WeaviateClient( - embedded_options=EmbeddedOptions( - persistence_data_path=data_path, - additional_env_vars=additional_env_vars, - ), - ) - return client + conn = sqlite3.connect(self.db_path) + conn.enable_load_extension(True) + sqlite_vec.load(conn) + conn.enable_load_extension(False) + return conn except Exception as e: - logger.error(f"Error during client creation: {str(e)}") - return None - - def setup_schema(self, client): - for class_config in self.schema_config: - if not client.collections.exists(class_config["name"]): - client.collections.create( - class_config["name"], properties=class_config["properties"] - ) - logger.info(f"Weaviate schema for class {class_config['name']} setup complete.") + logger.error("Failed to initialize database connection", error=str(e)) + raise + + 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 + ) + """ + ) - async def search_by_property(self, name: str, properties: List[str]) -> list[object]: - if len(properties) == 0: - return [] + # 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)") - # Perform the vector search - if self.weaviate_client is None: - logger.error("Could not find client, not returning results.") - return [] + self.conn.commit() - if not self.weaviate_client: - logger.error("Invalid client, cannot perform search.") + async def search_by_property(self, name: str, properties: List[str]) -> list[object]: + if len(properties) == 0: return [] try: - packages = self.weaviate_client.collections.get("Package") - response = packages.query.fetch_objects( - filters=Filter.by_property(name).contains_any(properties), - ) - - if not response: - return [] - - # Weaviate performs substring matching of the properties. So - # we need to double check the response. - properties = [prop.lower() for prop in properties] - filterd_objects = [] - for object in response.objects: - if object.properties[name].lower() in properties: - filterd_objects.append(object) - response.objects = filterd_objects - - return response.objects + cursor = self.conn.cursor() + placeholders = ",".join("?" * len(properties)) + query = f""" + SELECT name, type, status, description + FROM packages + WHERE LOWER({name}) IN ({placeholders}) + """ # nosec + + cursor.execute(query, [prop.lower() for prop in properties]) + results = [] + for row in cursor.fetchall(): + results.append( + { + "properties": { + "name": row[0], + "type": row[1], + "status": row[2], + "description": row[3], + } + } + ) + return results except Exception as e: - logger.error(f"An error occurred: {str(e)}") + logger.error(f"An error occurred during property search: {str(e)}") return [] async def search( @@ -152,66 +130,99 @@ async def search( distance: float = 0.3, ) -> list[object]: """ - Search the 'Package' collection based on a query string, ecosystem and packages. - If packages and ecosystem are both not none, then filter the objects using them. - If packages is not none and ecosystem is none, then filter the objects using - package names. - If packages is none, then perform vector search. - - Args: - query (str): The text query for which to search. - limit (int): The number of results to return. - distance (float): The distance threshold for the search. - ecosystem (str): The ecosystem to search in. - packages (list): The list of packages to filter the search. - - Returns: - list: A list of matching results with their properties and distances. + Search packages based on vector similarity or direct property matches. """ - try: - collection = self.weaviate_client.collections.get("Package") + cursor = self.conn.cursor() - response = None if packages and ecosystem and ecosystem in VALID_ECOSYSTEMS: - response = collection.query.fetch_objects( - filters=wvc.query.Filter.all_of([ - wvc.query.Filter.by_property("name").contains_any(packages), - wvc.query.Filter.by_property("type").equal(ecosystem) - ]), + placeholders = ",".join("?" * len(packages)) + query_sql = f""" + SELECT name, type, status, description + FROM packages + WHERE LOWER(name) IN ({placeholders}) + AND LOWER(type) = ? + """ # nosec + params = [p.lower() for p in packages] + [ecosystem.lower()] + logger.debug( + "Searching by package names and ecosystem", + packages=packages, + ecosystem=ecosystem, + sql=query_sql, + params=params, ) - response.objects = [ - obj - for obj in response.objects - if obj.properties["name"].lower() in packages - and obj.properties["type"].lower() == ecosystem.lower() - ] + cursor.execute(query_sql, params) + elif packages and not ecosystem: - response = collection.query.fetch_objects( - filters=wvc.query.Filter.all_of([ - wvc.query.Filter.by_property("name").contains_any(packages), - ]), + placeholders = ",".join("?" * len(packages)) + query_sql = f""" + SELECT name, type, status, description + FROM packages + WHERE LOWER(name) IN ({placeholders}) + """ # nosec + params = [p.lower() for p in packages] + logger.debug( + "Searching by package names only", + packages=packages, + sql=query_sql, + params=params, ) - response.objects = [ - obj - for obj in response.objects - if obj.properties["name"].lower() in packages - ] + cursor.execute(query_sql, params) + elif query: - # Perform the vector search - # Generate the vector for the query + # Generate embedding for the query query_vector = await self.inference_engine.embed(self.model_path, [query]) - - response = collection.query.near_vector( - query_vector[0], + query_embedding = np.array(query_vector[0], dtype=np.float32) + query_embedding_bytes = query_embedding.tobytes() + + query_sql = """ + WITH distances AS ( + SELECT name, type, status, description, + vec_distance_cosine(embedding, ?) as distance + FROM packages + ) + SELECT name, type, status, description, distance + FROM distances + WHERE distance <= ? + ORDER BY distance ASC + LIMIT ? + """ # nosec + logger.debug( + "Performing vector similarity search", + query=query, + distance_threshold=distance, limit=limit, - distance=distance, - return_metadata=MetadataQuery(distance=True), ) - - if not response: + cursor.execute(query_sql, (query_embedding_bytes, distance, limit)) + else: return [] - return response.objects + + # Log the raw SQL results + rows = cursor.fetchall() + logger.debug( + "Raw SQL results", + row_count=len(rows), + rows=[ + {"name": row[0], "type": row[1], "status": row[2], "description": row[3]} + for row in rows + ], + ) + + results = [] + for row in rows: + result = { + "properties": { + "name": row[0], + "type": row[1], + "status": row[2], + "description": row[3], + } + } + if query: # Add distance for vector searches + result["metadata"] = {"distance": row[4]} + results.append(result) + + return results except Exception as e: logger.error(f"Error during search: {str(e)}") diff --git a/src/codegate/storage/utils.py b/src/codegate/storage/utils.py index 45b77a85..f9140fe8 100644 --- a/src/codegate/storage/utils.py +++ b/src/codegate/storage/utils.py @@ -1,22 +1,29 @@ -import weaviate -from weaviate.embedded import EmbeddedOptions +import os +import shutil def restore_storage_backup(backup_path: str, backup_name: str): - client = weaviate.WeaviateClient( - embedded_options=EmbeddedOptions( - persistence_data_path="./weaviate_data", - grpc_port=50052, - additional_env_vars={ - "ENABLE_MODULES": "backup-filesystem", - "BACKUP_FILESYSTEM_PATH": backup_path, - }, - ) - ) - client.connect() + """ + Restore a SQLite database backup. + + Args: + backup_path: The directory containing the backup + backup_name: The name of the backup file + """ + backup_file = os.path.join(backup_path, backup_name) + target_dir = "./sqlite_data" + target_file = os.path.join(target_dir, "vectordb.db") + + if not os.path.exists(backup_file): + print(f"Backup file not found: {backup_file}") + return + try: - client.backup.restore(backup_id=backup_name, backend="filesystem", wait_for_completion=True) + # Create target directory if it doesn't exist + os.makedirs(target_dir, exist_ok=True) + + # Copy the backup file to the target location + shutil.copy2(backup_file, target_file) + print(f"Successfully restored backup from {backup_file}") except Exception as e: print(f"Failed to restore backup: {e}") - finally: - client.close() diff --git a/tests/test_storage.py b/tests/test_storage.py deleted file mode 100644 index dafb1547..00000000 --- a/tests/test_storage.py +++ /dev/null @@ -1,60 +0,0 @@ -from unittest.mock import AsyncMock, MagicMock, patch - -import pytest - -from codegate.config import Config -from codegate.storage.storage_engine import ( - StorageEngine, -) # Adjust the import based on your actual path - - -@pytest.fixture -def mock_weaviate_client(): - client = MagicMock() - response = MagicMock() - response.objects = [ - { - "properties": { - "name": "test", - "type": "library", - "status": "active", - "description": "test description", - } - } - ] - client.collections.get.return_value.query.near_vector.return_value = response - return client - - -@pytest.fixture -def mock_inference_engine(): - engine = AsyncMock() - engine.embed.return_value = [0.1, 0.2, 0.3] # Adjust based on expected vector dimensions - return engine - - -@pytest.mark.asyncio -async def test_search(mock_weaviate_client, mock_inference_engine): - Config.load(config_path="./config.yaml") - - # Patch the LlamaCppInferenceEngine.embed method (not the entire class) - with patch( - "codegate.inference.inference_engine.LlamaCppInferenceEngine.embed", - mock_inference_engine.embed, - ): - # Initialize StorageEngine - with patch( - "codegate.storage.storage_engine.StorageEngine.get_client", - return_value=mock_weaviate_client, - ): - # Initialize StorageEngine - # Need to recreate instance to use the mock - storage_engine = StorageEngine.recreate_instance(data_path="./weaviate_data") - - # Invoke the search method - results = await storage_engine.search(query="test query", limit=5, distance=0.3) - - # Assertions to validate the expected behavior - assert len(results) == 1 # Assert that one result is returned - assert results[0]["properties"]["name"] == "test" - mock_weaviate_client.connect.assert_called() diff --git a/tests/vectordb/test_sqlitevec.py b/tests/vectordb/test_sqlitevec.py new file mode 100644 index 00000000..f80b644d --- /dev/null +++ b/tests/vectordb/test_sqlitevec.py @@ -0,0 +1,153 @@ +import os +from unittest.mock import AsyncMock, patch + +import numpy as np +import pytest + +from codegate.config import Config +from codegate.storage.storage_engine import StorageEngine + + +@pytest.fixture(scope="module") +def mock_sqlite_vec(): + with patch("sqlite_vec.load") as mock_load: + # Mock the vector similarity extension loading + def setup_vector_similarity(conn): + cursor = conn.cursor() + # Create a table to store our mock distance function + cursor.execute( + """ + CREATE TABLE IF NOT EXISTS vector_distances ( + id INTEGER PRIMARY KEY, + distance FLOAT + )""" + ) + # Insert a mock distance value that will be used in searches + cursor.execute("INSERT INTO vector_distances (distance) VALUES (0.1)") + + # Create a view that simulates the vec_distance_cosine function + cursor.execute( + """ + CREATE VIEW IF NOT EXISTS vec_distance_cosine_view AS + SELECT distance FROM vector_distances WHERE id = 1 + """ + ) + + # Create a function that returns the mock distance + conn.create_function("vec_distance_cosine", 2, lambda x, y: 0.1) + + mock_load.side_effect = setup_vector_similarity + yield mock_load + + +@pytest.fixture(scope="module") +def test_db_path(): + return "./test_sqlite_data/vectordb.db" + + +@pytest.fixture(scope="module") +def mock_config(test_db_path): + # Create a mock config instance + config = Config() + config.model_base_path = "./codegate_volume/models" + config.embedding_model = "all-minilm-L6-v2-q5_k_m.gguf" + config.vec_db_path = test_db_path + + # Mock Config.get_config to return our test config + with patch("codegate.config.Config.get_config", return_value=config): + yield config + + +@pytest.fixture(scope="module") +def storage_engine(mock_sqlite_vec, mock_config, test_db_path): + # Setup: Create a temporary database for testing + test_db_dir = os.path.dirname(test_db_path) + os.makedirs(test_db_dir, exist_ok=True) + + engine = StorageEngine(data_path=test_db_dir) + yield engine + # Teardown: Remove the temporary database and directory + del engine + if os.path.exists(test_db_path): + os.remove(test_db_path) + os.rmdir(test_db_dir) + + +@pytest.fixture(autouse=True) +def clean_database(storage_engine): + # Clear all data before each test + cursor = storage_engine.conn.cursor() + cursor.execute("DELETE FROM packages") + storage_engine.conn.commit() + yield + + +@pytest.mark.asyncio +async def test_search_by_property(storage_engine): + # Insert test data + cursor = storage_engine.conn.cursor() + cursor.execute( + """ + INSERT INTO packages (name, type, status, description) + VALUES ('invokehttp', 'pypi', 'active', 'An evil package') + """ + ) + storage_engine.conn.commit() + + # Test search by property + results = await storage_engine.search_by_property("name", ["invokehttp"]) + assert len(results) == 1 + assert results[0]["properties"]["name"] == "invokehttp" + assert results[0]["properties"]["type"] == "pypi" + assert results[0]["properties"]["status"] == "active" + assert results[0]["properties"]["description"] == "An evil package" + + +@pytest.mark.asyncio +async def test_search_by_package_names(storage_engine): + # Insert test data + cursor = storage_engine.conn.cursor() + cursor.execute( + """ + INSERT INTO packages (name, type, status, description) + VALUES ('invokehttp', 'pypi', 'active', 'An evil package') + """ + ) + storage_engine.conn.commit() + + # Test search by package names + results = await storage_engine.search(packages=["invokehttp"]) + assert len(results) == 1 + assert results[0]["properties"]["name"] == "invokehttp" + assert results[0]["properties"]["type"] == "pypi" + assert results[0]["properties"]["status"] == "active" + assert results[0]["properties"]["description"] == "An evil package" + + +@pytest.mark.asyncio +async def test_search_by_query(storage_engine): + # Mock the inference engine to return a fixed embedding + with patch.object( + storage_engine.inference_engine, "embed", new=AsyncMock(return_value=[[0.1, 0.2, 0.3]]) + ): + # Insert test data with embedding + cursor = storage_engine.conn.cursor() + embedding = np.array([0.1, 0.2, 0.3], dtype=np.float32).tobytes() + cursor.execute( + """ + INSERT INTO packages (name, type, status, description, embedding) + VALUES ('invokehttp', 'pypi', 'active', 'An evil package', ?) + """, + (embedding,), + ) + storage_engine.conn.commit() + + # Test search by query + results = await storage_engine.search(query="test") + assert len(results) == 1 + assert results[0]["properties"]["name"] == "invokehttp" + assert results[0]["properties"]["type"] == "pypi" + assert results[0]["properties"]["status"] == "active" + assert results[0]["properties"]["description"] == "An evil package" + assert "metadata" in results[0] + assert "distance" in results[0]["metadata"]