Updates

Author: ParamThakkar123

src/JuliaHealthLLM.jl

@@ -1,25 +1,19 @@
 module JuliaHealthLLM
 
-using LibGit2
-using Downloads
-using Dates
 using PromptingTools
-using TextAnalysis
-using Serialization, Base.Threads
-using Transformers, CUDA
-using LlamaCpp, JSON3, HuggingFaceTokenizers, Jjama3
-using JSON, Markdown
+using PromptingTools.Experimental.RAGTools
+using LinearAlgebra, SparseArrays
+using JSON3, Serialization
+using Statistics
+using LibPQ
 
-# TODO: Add functions from scripts
-include("../scripts/clone.jl")
-include("../scripts/knowledge.jl")
-include("../scripts/restore_submodules.jl")
-include("../scripts/update_julia_package_repos.jl")
-include("../scripts/update_and_record.jl")
-include("../src/chunk_docs.jl")
+include("utils.jl")
+include("pgvector.jl")
+include("database.jl")
+include("embedding.jl")
+include("query.jl")
 
-export clone_repositories, normalize_repo_url, clone_repo, is_textual_file,
-       build_corpus_from_repo, save_corpus, process_directory, update_and_record_submodules,
-       restore_all_submodules, update_julia_package_repos, chunk_file_content, traverse_repository
+export collect_files_with_extensions, write_combined_file, generate_funsql_query,
+build_index_rag, store_embeddings_pgvector
 
 end

src/chunk_docs.jl

@@ -0,0 +1,28 @@
+module Database
+
+function store_embeddings_pgvector(conn::LibPQ.Connection, embeddings::AbstractMatrix, chunks::AbstractVector, embedding_dimension::Int)
+    LibPQ.execute(conn, """
+        CREATE TABLE IF NOT EXISTS embeddings (
+            id SERIAL PRIMARY KEY,
+            chunk TEXT NOT NULL,
+            embedding VECTOR($embedding_dimension)
+        )
+    """)
+
+    embeddings = convert(Matrix{Float64}, embeddings)
+
+    chunks = String.(chunks)
+
+    for i in 1:eachindex(embeddings, 2)
+        chunk = chunks[i]
+        embedding = Pgvector.convert(embeddings[:, i])
+        LibPQ.execute(conn, """
+            INSERT INTO embeddings (chunk, embedding)
+            VALUES (\$1, \$2)
+        """, (chunk, embedding))
+    end
+
+    println("Embeddings successfully stored in the database.")
+end
+
+end

src/database.jl

@@ -0,0 +1,7 @@
+module Embedding
+
+function build_index_rag(cfg, files; embedder_kwargs=())
+    return PromptingTools.Experimental.RAGTools.build_index(cfg, files; embedder_kwargs=embedder_kwargs)
+end
+
+end

src/dummy_src_file.jl

@@ -0,0 +1,3 @@
+module Pgvector
+    convert(v::AbstractVector{T}) where T<:Real = string("[", join(v, ","), "]")
+end

src/embedding.jl

@@ -0,0 +1,21 @@
+module Query
+
+function generate_funsql_query(index, model_embedding::String, model_name::String, prompt_template::String, question::String)
+    PromptingTools.MODEL_CHAT = model_name
+    PromptingTools.MODEL_EMBEDDING = model_embedding
+
+    PT.register_model!(name=MODEL_CHAT, schema=PT.OllamaSchema())
+    PT.register_model!(name=MODEL_EMBEDDING, schema=PT.OllamaSchema())
+    
+    prompt = replace(prompt_template, "{input_query}" => question)
+
+    answer = airag(index; 
+        question=prompt,
+        retriever_kwargs=(model=model_embedding, schema=PT.OllamaSchema(), embedder_kwargs=(schema=PT.OllamaSchema(), model=MODEL_EMBEDDING)),
+        generator_kwargs=(model=model_name, schema=PT.OllamaSchema())
+    )
+
+    return answer
+end
+
+end

src/pgvector.jl

@@ -0,0 +1,30 @@
+module Utils
+
+function collect_files_with_extensions(directory::String, extensions::Vector{String})
+    files = String[]
+    for (root, _, file_names) in walkdir(directory)
+        for file_name in file_names
+            if any(endswith(file_name, ext) for ext in extensions)
+                relative_path = joinpath(relpath(root, directory), file_name)
+                push!(files, relative_path)
+            end
+        end
+    end
+    return files
+end
+
+function write_combined_file(files::Vector{String}, output_file::String)
+    open(output_file, "w") do io
+        for file in files
+            println(io, "# File: $file")
+            open(file, "r") do f
+                for line in eachline(f)
+                    println(io, line)
+                end
+            end
+            println(io, "\n")
+        end
+    end
+end
+
+end