corpus-engine-tanjha 0.1.0__tar.gz

corpus_engine_tanjha-0.1.0/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2026 Tanush Ojha
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

corpus_engine_tanjha-0.1.0/PKG-INFO

@@ -0,0 +1,34 @@
+Metadata-Version: 2.4
+Name: corpus-engine_tanjha
+Version: 0.1.0
+Summary: A CLI tool meant for directory text querying VIA text embedding
+Author: Tanush Ojha
+Author-email: Tanush Ojha <tanushojha057@gmail.com>
+License-Expression: MIT
+License-File: LICENSE
+Requires-Dist: einops>=0.8.2
+Requires-Dist: langchain-text-splitters>=1.1.2
+Requires-Dist: numpy>=2.4.6
+Requires-Dist: scikit-learn>=1.9.0
+Requires-Dist: sentence-transformers>=5.5.1
+Requires-Dist: typer>=0.26.5
+Requires-Python: >=3.11.9
+Project-URL: Homepage, https://github.com/tanjha/corpus-engine
+Project-URL: Issues, https://github.com/tanjha/corpus-engine/issues
+Description-Content-Type: text/markdown
+
+# Corpus Engine
+Vector Embedding Powered Corpus Search Engine
+
+
+# Build
+
+Build in project root using:
+```
+uv tool install .
+```
+
+Dev build using:
+```
+uv tool install -e .
+```

corpus_engine_tanjha-0.1.0/README.md

@@ -0,0 +1,15 @@
+# Corpus Engine
+Vector Embedding Powered Corpus Search Engine
+
+
+# Build
+
+Build in project root using:
+```
+uv tool install .
+```
+
+Dev build using:
+```
+uv tool install -e .
+```

corpus_engine_tanjha-0.1.0/pyproject.toml

@@ -0,0 +1,32 @@
+[project]
+name = "corpus-engine_tanjha"
+version = "0.1.0"
+description = "A CLI tool meant for directory text querying VIA text embedding"
+authors = [{name="Tanush Ojha",email="tanushojha057@gmail.com"}]
+readme = "README.md"
+requires-python = ">=3.11.9"
+dependencies = [
+    "einops>=0.8.2",
+    "langchain-text-splitters>=1.1.2",
+    "numpy>=2.4.6",
+    "scikit-learn>=1.9.0",
+    "sentence-transformers>=5.5.1",
+    "typer>=0.26.5",
+]
+license= "MIT"
+license-files=["LICEN[CS]E*"]
+
+[project.scripts]
+engine = "corpusEngine.engine:app"
+
+[project.urls]
+Homepage = "https://github.com/tanjha/corpus-engine"
+Issues = "https://github.com/tanjha/corpus-engine/issues"
+
+
+[build-system]
+requires = ["uv_build >= 0.11.23, <0.12.0"]
+build-backend = "uv_build"
+
+[tool.uv.build-backend]
+module-name = "corpusEngine"

corpus_engine_tanjha-0.1.0/src/corpusEngine/__main__.py

@@ -0,0 +1,4 @@
+if __name__ == "__main__":
+    from corpusEngine.cli import app
+
+    app()

corpus_engine_tanjha-0.1.0/src/corpusEngine/dbManager.py

@@ -0,0 +1,134 @@
+import sqlite3
+from pathlib import Path
+import numpy as np
+from .entryObj import entry
+
+
+class dbManager:
+    def __init__(self, db: sqlite3.Connection):
+        self.db = db
+
+    def initSQL(self):
+        self.db.execute("""
+            CREATE TABLE IF NOT EXISTS files (
+                id      INTEGER PRIMARY KEY,
+                path         TEXT UNIQUE NOT NULL,
+                time_updated TEXT NOT NULL
+            );
+        """)
+        self.db.execute("""
+            CREATE TABLE IF NOT EXISTS chunks (
+                id         INTEGER PRIMARY KEY,
+                file_id    INTEGER NOT NULL REFERENCES files(id) ON DELETE CASCADE,
+                start_line INTEGER,
+                end_line   INTEGER,
+                chunk_num  INTEGER,
+                embedding  BLOB NOT NULL
+            );
+        """)
+        self.db.execute(
+            "CREATE INDEX IF NOT EXISTS idx_chunks_file_id ON chunks(file_id);"
+        )
+        self.db.commit()
+
+    def getEmbeddings(self):
+        rows = self.db.execute("""
+            SELECT c.id, c.start_line, c.end_line, f.path, c.embedding, c.chunk_num
+            FROM chunks c
+            JOIN files f ON c.file_id = f.id
+        """)
+        return [
+            {
+                "id": row[0],
+                "start_line": row[1],
+                "end_line": row[2],
+                "path": Path(row[3]),
+                "embedding": np.frombuffer(row[4], dtype=np.float32),
+                "chunk_num": int(row[5]),
+            }
+            for row in rows
+        ]
+
+    def getDB(self) -> sqlite3.Connection:
+        return self.db
+
+    def insertFull(self, entries: list[entry]):
+        if entries:
+            self._insertFile(str(entries[0].path.absolute()), entries[0].time)
+            print(f"inserting {str(entries[0].path.absolute())}")
+            self.db.commit()
+            for ent in entries:
+                file_id = self._getFileId(str(ent.path.absolute()))
+                self._insertChunk(
+                    file_id,
+                    ent.start_line,
+                    ent.end_line,
+                    ent.chunk_num,
+                    ent.embedding,
+                )
+            self.db.commit()
+
+    def removeAll(self, paths: list[Path]):
+        for path in paths:
+            p = str(path)
+            fileId = self._getFileId(path)
+            if fileId is not None:
+                self.db.execute("DELETE FROM files where path = ?", (p,))
+        self.db.commit()
+
+    def removeFile(self, path: Path):
+        p = str(path)
+        fileId = self._getFileId(path)
+        if fileId is not None:
+            self.db.execute("DELETE FROM files where path = ?", (p,))
+            self.db.commit()
+
+    def _insertFile(
+        self,
+        path: Path,
+        time: str,
+    ):
+        self.db.execute(
+            """INSERT INTO files (path, time_updated)
+            VALUES (?, ?)""",
+            (
+                path,
+                time,
+            ),
+        )
+
+    def _getFileId(self, path: Path) -> int:
+        p = str(path)
+        row = self.db.execute("SELECT id FROM files WHERE path = ?", (p,)).fetchone()
+        if row is None:
+            return row
+        return row[0]
+
+    def getFileTime(self, path: Path):
+        p = str(path)
+        row = self.db.execute(
+            "SELECT time_updated FROM files where path = ?", (p,)
+        ).fetchone()
+        if row is None:
+            return row
+        return row[0]
+
+    def _insertChunk(
+        self,
+        file_id: int,
+        start_line: int,
+        end_line: int,
+        chunk_num: int,
+        embedding: np.ndarray,
+    ):
+        self.db.execute(
+            """INSERT INTO chunks (file_id, start_line, end_line, chunk_num, embedding)
+            VALUES (?, ?, ?, ?, ?)""",
+            (
+                file_id,
+                start_line,
+                end_line,
+                chunk_num,
+                embedding.astype(np.float32).tobytes(),
+            ),
+        )

corpus_engine_tanjha-0.1.0/src/corpusEngine/embed.py

@@ -0,0 +1,42 @@
+from sentence_transformers import SentenceTransformer, util
+
+
+class embed_engine:
+    def __init__(self, embeddings):
+        self.model = SentenceTransformer(
+            "nomic-ai/nomic-embed-text-v1", trust_remote_code=True
+        )
+        self.embeddings = embeddings
+
+    def embedDoc(self, raw_text):
+        embedding = self.model.encode(
+            f"search_document: {raw_text}", normalize_embeddings=True
+        )
+        return embedding
+
+    def embedAll(self, texts: list[str]):
+        prefixed = [f"search_document: {text}" for text in texts]
+        embedding = self.model.encode(
+            prefixed, normalize_embeddings=True, batch_size=32, show_progress_bar=True
+        )
+        return embedding
+
+    def _embedQuery(self, raw_text):
+        embedding = self.model.encode(
+            f"search_query: {raw_text}", normalize_embeddings=True
+        )
+        return embedding
+
+    def similarityFull(self, query, embeddings):
+        for embed in embeddings:
+            embed["similarity"] = self._similarity(query, embed["embedding"])
+
+        return embeddings
+
+    def _similarity(self, query, embed):
+        similarity = util.cos_sim(query, embed)
+        return similarity
+
+    def setEmbed(self, embeddings):
+        self.embeddings = embeddings
+        return

corpus_engine_tanjha-0.1.0/src/corpusEngine/engine.py

@@ -0,0 +1,81 @@
+import typer
+from pathlib import Path
+from typing import Annotated
+import os
+
+app = typer.Typer(no_args_is_help=True)
+
+
+@app.command()
+def init(debug: Annotated[bool, typer.Option(help="Print debug logs")] = False):
+    """Initialize engine into current directory."""
+    doInit(debug)
+
+
+def doInit(debug: bool = False):
+    initialized = False
+    cwd = Path.cwd()
+    # print(cwd)
+    for item in cwd.iterdir():
+        if item.name == "corpusEngine.db":
+            initialized = True
+            print("[INFO] Corpus Engine already initialized")
+    if not initialized:
+        from .worker import setup
+
+        setup(cwd, debug)
+
+
+@app.command()
+def query(debug: Annotated[bool, typer.Option(help="Print debug logs")] = False):
+    """Query directory using input - returns top 5 search results"""
+    initialized = False
+    cwd = Path.cwd()
+    # print(cwd)
+    for item in cwd.iterdir():
+        if item.name == "corpusEngine.db":
+            initialized = True
+
+    if initialized:
+        from .worker import query
+
+        query(debug)
+    else:
+        print("[INFO] Corpus Engine not initialized")
+
+
+@app.command()
+def update():
+    initialized = False
+    cwd = Path.cwd()
+    # print(cwd)
+    for item in cwd.iterdir():
+        if item.name == "corpusEngine.db":
+            initialized = True
+
+    if initialized:
+        from .worker import updateEmbed
+
+        updateEmbed(cwd)
+    else:
+        print("[INFO] Corpus Engine not initialized")
+
+
+@app.command()
+def reset(debug: Annotated[bool, typer.Option(help="Print debug logs")] = False):
+    """Remove Corpus Engine and reinitialize it."""
+    doRemove()
+    doInit(debug)
+
+
+@app.command()
+def remove():
+    doRemove()
+
+
+def doRemove():
+    if os.path.exists("corpusEngine.db"):
+        os.remove("corpusEngine.db")
+        print("[INFO] Removed Corpus Engine")
+    else:
+        print("[INFO] Corpus Engine not initialized")

corpus_engine_tanjha-0.1.0/src/corpusEngine/entryObj.py

@@ -0,0 +1,35 @@
+from pathlib import Path
+import numpy as np
+
+
+class entry:
+    path: Path
+    start_line: int
+    end_line: int
+    time: str
+    chunk_num: int
+    embedding: np.ndarray
+
+    def __init__(
+        self,
+        path: Path,
+        start_line: int,
+        end_line: int,
+        time: str,
+        chunk_num: int,
+    ):
+        self.path = path
+        self.start_line = start_line
+        self.end_line = end_line
+        self.time = time
+        self.chunk_num = chunk_num
+        self.embedding = None
+
+    def set_embed(self, embed: np.ndarray):
+        self.embedding = embed
+        return
+
+    def printAll(self):
+        print(
+            f"Path: {self.path} Start: {self.start_line} End: {self.end_line} Time {self.time} Raw: {self.chunk_num}"
+        )

corpus_engine_tanjha-0.1.0/src/corpusEngine/vectorCheck.py

@@ -0,0 +1,464 @@
+"""
+vectorizable_extensions.py
+==========================
+
+A comprehensive catalog of file-type suffixes whose contents can be turned into
+human-readable plaintext, and therefore chunked + embedded ("vectorized") for
+search / RAG pipelines.
+
+Two practical buckets:
+
+  1. PLAINTEXT_NATIVE  -> the bytes on disk are already text; just decode them.
+  2. NEEDS_EXTRACTION  -> binary/container formats that yield text only after a
+                          parser/library extracts it (PDF, DOCX, EPUB, ...).
+
+Use `VECTORIZABLE_EXTENSIONS` for a single flat membership check, or the
+category dicts below if you want to route each file to the right loader.
+
+All suffixes are lowercase and include the leading dot.
+
+Optional buckets (images via OCR, audio/video via transcription) are provided
+separately at the bottom and are NOT included in the defaults, because their
+"text" is derived rather than contained in the file.
+"""
+
+from pathlib import Path
+
+# ---------------------------------------------------------------------------
+# 1. PLAINTEXT-NATIVE  (read directly, e.g. open(path, encoding="utf-8"))
+# ---------------------------------------------------------------------------
+
+PLAIN_TEXT = {
+    ".txt",
+    ".text",
+    ".log",
+    ".nfo",
+    ".me",
+    ".readme",
+    ".asc",
+    ".diff",
+    ".patch",
+}
+
+MARKUP_AND_DOCS_TEXT = {
+    ".md",
+    ".markdown",
+    ".mdown",
+    ".mkd",
+    ".mdx",
+    ".rst",
+    ".adoc",
+    ".asciidoc",
+    ".org",
+    ".textile",
+    ".creole",
+    ".wiki",
+    ".tex",
+    ".latex",
+    ".ltx",
+    ".bib",
+    ".sty",
+    ".cls",
+    ".man",
+    ".roff",
+    ".1",
+    ".2",
+    ".3",
+}
+
+WEB_AND_TEMPLATES = {
+    ".html",
+    ".htm",
+    ".xhtml",
+    ".shtml",
+    ".css",
+    ".scss",
+    ".sass",
+    ".less",
+    ".styl",
+    ".vue",
+    ".svelte",
+    ".astro",
+    ".jsx",
+    ".tsx",
+    ".haml",
+    ".pug",
+    ".jade",
+    ".ejs",
+    ".erb",
+    ".hbs",
+    ".handlebars",
+    ".mustache",
+    ".liquid",
+    ".twig",
+    ".njk",
+    ".j2",
+    ".jinja",
+    ".jinja2",
+}
+
+DATA_AND_CONFIG = {
+    ".json",
+    ".jsonl",
+    ".ndjson",
+    ".json5",
+    ".jsonc",
+    ".geojson",
+    ".yaml",
+    ".yml",
+    ".toml",
+    ".ini",
+    ".cfg",
+    ".conf",
+    ".config",
+    ".properties",
+    ".env",
+    ".csv",
+    ".tsv",
+    ".psv",
+    ".tab",
+    ".xml",
+    ".xsd",
+    ".xsl",
+    ".xslt",
+    ".dtd",
+    ".rss",
+    ".atom",
+    ".plist",
+    ".svg",
+    ".graphql",
+    ".gql",
+    ".proto",
+    ".thrift",
+    ".avsc",
+    ".prisma",
+    ".ron",
+    ".hcl",
+    ".editorconfig",
+    ".gitignore",
+    ".gitattributes",
+    ".dockerignore",
+    ".npmrc",
+    ".yarnrc",
+}
+
+SOURCE_CODE = {
+    # Python
+    ".py",
+    ".pyw",
+    ".pyi",
+    ".pyx",
+    ".pxd",
+    ".rpy",
+    # JavaScript / TypeScript
+    ".js",
+    ".mjs",
+    ".cjs",
+    ".ts",
+    ".mts",
+    ".cts",
+    ".coffee",
+    # C / C++
+    ".c",
+    ".h",
+    ".cpp",
+    ".cc",
+    ".cxx",
+    ".c++",
+    ".hpp",
+    ".hh",
+    ".hxx",
+    ".h++",
+    ".ino",
+    # C# / .NET
+    ".cs",
+    ".csx",
+    ".vb",
+    ".fs",
+    ".fsx",
+    ".fsi",
+    # JVM languages
+    ".java",
+    ".kt",
+    ".kts",
+    ".scala",
+    ".sc",
+    ".groovy",
+    ".gradle",
+    ".clj",
+    ".cljs",
+    ".cljc",
+    ".edn",
+    # Go / Rust / Zig / Nim / D / Crystal
+    ".go",
+    ".rs",
+    ".zig",
+    ".nim",
+    ".nims",
+    ".d",
+    ".cr",
+    # Ruby
+    ".rb",
+    ".rake",
+    ".gemspec",
+    ".ru",
+    # PHP / Perl
+    ".php",
+    ".phtml",
+    ".php3",
+    ".php4",
+    ".php5",
+    ".pl",
+    ".pm",
+    ".pod",
+    ".t",
+    # Swift / Objective-C
+    ".swift",
+    ".m",
+    ".mm",
+    # Shell / scripting
+    ".sh",
+    ".bash",
+    ".zsh",
+    ".fish",
+    ".ksh",
+    ".csh",
+    ".tcsh",
+    ".ps1",
+    ".psm1",
+    ".psd1",
+    ".bat",
+    ".cmd",
+    # Lua / Tcl / Dart / Julia / R
+    ".lua",
+    ".tcl",
+    ".dart",
+    ".jl",
+    ".r",
+    ".rmd",
+    ".rnw",
+    # Functional / academic
+    ".ex",
+    ".exs",
+    ".erl",
+    ".hrl",
+    ".hs",
+    ".lhs",
+    ".ml",
+    ".mli",
+    ".elm",
+    ".scm",
+    ".ss",
+    ".lisp",
+    ".lsp",
+    ".cl",
+    ".el",
+    ".rkt",
+    # Systems / legacy
+    ".pas",
+    ".pp",
+    ".f",
+    ".for",
+    ".f90",
+    ".f95",
+    ".f03",
+    ".f08",
+    ".cob",
+    ".cbl",
+    ".asm",
+    ".s",
+    ".vhd",
+    ".vhdl",
+    ".v",
+    ".sv",
+    ".vbs",
+    # Other
+    ".sql",
+    ".sol",
+    ".gd",
+    ".gdscript",
+    ".applescript",
+    ".scpt",
+    ".vala",
+    ".hx",
+    ".moon",
+    ".p",
+    ".ada",
+    ".adb",
+    ".ads",
+}
+
+BUILD_AND_INFRA = {
+    ".dockerfile",
+    ".containerfile",
+    ".mk",
+    ".cmake",
+    ".make",
+    ".bazel",
+    ".bzl",
+    ".buck",
+    ".starlark",
+    ".star",
+    ".tf",
+    ".tfvars",
+    ".nomad",
+    ".jenkinsfile",
+    ".cabal",
+    ".lock",
+}
+
+NOTEBOOKS_AND_LITERATE = {
+    ".ipynb",  # JSON container; cell text is readable after parsing
+    ".rmd",  # already above, kept for clarity
+}
+
+SUBTITLES = {
+    ".srt",
+    ".vtt",
+    ".sub",
+    ".ass",
+    ".ssa",
+    ".sbv",
+    ".lrc",
+}
+
+EMAIL_TEXT = {
+    ".eml",
+    ".mbox",
+    ".mbx",
+}
+
+# Aggregate of everything that is already plaintext on disk
+PLAINTEXT_NATIVE = (
+    PLAIN_TEXT
+    | MARKUP_AND_DOCS_TEXT
+    | WEB_AND_TEMPLATES
+    | DATA_AND_CONFIG
+    | SOURCE_CODE
+    | BUILD_AND_INFRA
+    | NOTEBOOKS_AND_LITERATE
+    | SUBTITLES
+    | EMAIL_TEXT
+)
+
+# ---------------------------------------------------------------------------
+# 2. NEEDS-EXTRACTION  (binary/container; text via a parser library)
+#    Suggested library noted per format.
+# ---------------------------------------------------------------------------
+
+NEEDS_EXTRACTION = {
+    # Documents
+    ".pdf",  # pdfplumber / pypdf / PyMuPDF (fitz) / unstructured
+    ".doc",  # antiword / textract / LibreOffice
+    ".docx",  # python-docx / docx2txt / unstructured
+    ".odt",  # odfpy
+    ".rtf",  # striprtf
+    ".wpd",  # libwpd
+    ".epub",  # ebooklib / unstructured
+    ".mobi",  # mobi / kindleunpack
+    ".azw",
+    ".azw3",
+    ".fb2",  # ebook-convert (Calibre)
+    ".pages",  # Apple Pages (zip container; LibreOffice / textutil)
+    # Spreadsheets
+    ".xls",  # xlrd
+    ".xlsx",
+    ".xlsm",  # openpyxl / pandas
+    ".ods",  # odfpy / pandas
+    ".numbers",  # numbers-parser
+    # Presentations
+    ".ppt",  # LibreOffice / textract
+    ".pptx",  # python-pptx
+    ".odp",  # odfpy
+    ".key",  # Apple Keynote (LibreOffice)
+    # Email containers
+    ".msg",  # extract-msg
+    ".pst",
+    ".ost",  # libpff / readpst
+}
+
+# ---------------------------------------------------------------------------
+# 3. FLAT SET — use this for a quick "is this file vectorizable?" check
+# ---------------------------------------------------------------------------
+
+VECTORIZABLE_EXTENSIONS = frozenset(PLAINTEXT_NATIVE | NEEDS_EXTRACTION)
+
+# Sorted list form, if you prefer a list literal:
+VECTORIZABLE_EXTENSIONS_LIST = sorted(VECTORIZABLE_EXTENSIONS)
+
+# ---------------------------------------------------------------------------
+# 4. OPTIONAL — derived text, NOT included by default
+#    (the file does not *contain* text; text is produced from it)
+# ---------------------------------------------------------------------------
+
+OCR_IMAGE_EXTENSIONS = {
+    ".png",
+    ".jpg",
+    ".jpeg",
+    ".tif",
+    ".tiff",
+    ".bmp",
+    ".gif",
+    ".webp",
+    # via pytesseract / easyocr / unstructured
+}
+
+TRANSCRIBABLE_AV_EXTENSIONS = {
+    ".mp3",
+    ".wav",
+    ".m4a",
+    ".flac",
+    ".ogg",
+    ".aac",
+    ".mp4",
+    ".mov",
+    ".mkv",
+    ".avi",
+    ".webm",
+    ".m4v",
+    # via whisper / faster-whisper
+}
+
+# ---------------------------------------------------------------------------
+# 5. Helpers
+# ---------------------------------------------------------------------------
+
+
+def is_vectorizable(path, include_ocr=False, include_av=False) -> bool:
+    """Return True if `path`'s suffix can be turned into text."""
+    ext = Path(path).suffix.lower()
+    if ext in VECTORIZABLE_EXTENSIONS:
+        return True
+    if include_ocr and ext in OCR_IMAGE_EXTENSIONS:
+        return True
+    if include_av and ext in TRANSCRIBABLE_AV_EXTENSIONS:
+        return True
+    # Common extension-less text files (Dockerfile, Makefile, etc.)
+    if ext == "" and Path(path).name.lower() in {
+        "dockerfile",
+        "makefile",
+        "rakefile",
+        "gemfile",
+        "procfile",
+        "jenkinsfile",
+        "vagrantfile",
+        "license",
+        "readme",
+        "authors",
+    }:
+        return True
+    return False
+
+
+def needs_extraction_library(path) -> bool:
+    """True if the file is binary and requires a parser to read its text."""
+    return Path(path).suffix.lower() in NEEDS_EXTRACTION
+
+
+if __name__ == "__main__":
+    # print(f"Plaintext-native suffixes : {len(PLAINTEXT_NATIVE)}")
+    # print(f"Needs-extraction suffixes : {len(NEEDS_EXTRACTION)}")
+    # print(f"Total vectorizable        : {len(VECTORIZABLE_EXTENSIONS)}")
+    # print()
+    # print(VECTORIZABLE_EXTENSIONS_LIST)
+    pass

corpus_engine_tanjha-0.1.0/src/corpusEngine/worker.py

@@ -0,0 +1,182 @@
+from langchain_text_splitters import RecursiveCharacterTextSplitter
+import sqlite3
+from pathlib import Path
+from datetime import datetime
+from .embed import embed_engine
+from .entryObj import entry
+from .dbManager import dbManager
+from .vectorCheck import is_vectorizable
+
+scale = 2
+config = {"chunk_size": 256 * scale, "overlap": 26 * scale}
+debug = False
+
+
+def setup(cwd: Path, dbg: bool = False):
+    # debug = dbg
+    print("[INFO] Initializing Corpus Engine...")
+    dbM = _initDB("corpusEngine.db")
+    dbM.initSQL()
+
+    splitter = RecursiveCharacterTextSplitter(
+        chunk_size=config["chunk_size"],
+        chunk_overlap=config["overlap"],
+        separators=[""],
+        strip_whitespace=False,
+    )
+    embedder = embed_engine([])
+
+    _lookDir(splitter=splitter, path=cwd, dbM=dbM, embedder=embedder)
+    print("[INFO] Finished initialization.")
+
+
+def updateEmbed(cwd: Path, dbg: bool = False):
+    dbM = _initDB("corpusEngine.db")
+    embedder = embed_engine([])
+    splitter = RecursiveCharacterTextSplitter(
+        chunk_size=config["chunk_size"],
+        chunk_overlap=config["overlap"],
+        separators=[""],
+        strip_whitespace=False,
+    )
+    walked_dir = _listAll(cwd)
+    paths = {str(p.absolute()) for p in walked_dir}
+    db_rows = {
+        row[0]: row[1]
+        for row in dbM.getDB().execute("SELECT path, time_updated FROM files")
+    }
+    for path in paths:
+        time = datetime.fromtimestamp(Path(path).stat().st_mtime)
+        if path not in db_rows:
+            _splitFile(file=Path(path), splitter=splitter, dbM=dbM, embedder=embedder)
+        elif time != db_rows[path]:
+            print("[INFO] Updating file")
+            dbM.removeFile(Path(path))
+            _splitFile(file=Path(path), splitter=splitter, dbM=dbM, embedder=embedder)
+    for removed in db_rows.keys() - paths:
+        dbM.removeFile(Path(removed))
+
+
+def query(debug: bool = False):
+    print("Querying")
+    dbM = _initDB("corpusEngine.db")
+    embeddings = dbM.getEmbeddings()
+    embedder = embed_engine(embeddings)
+    query = _getQuery(embedder)
+    sim_embeddings = embedder.similarityFull(query, embeddings)
+    similarities = sorted(sim_embeddings, key=lambda x: x["similarity"], reverse=True)[
+        :5
+    ]
+    splitter = RecursiveCharacterTextSplitter(
+        chunk_size=config["chunk_size"],
+        chunk_overlap=config["overlap"],
+        separators=[""],
+        strip_whitespace=False,
+    )
+
+    for sim in similarities:
+        similarity = sim["similarity"]
+        path = sim["path"]
+        start_line = sim["start_line"]
+        end_line = sim["end_line"]
+
+        chunk_num = sim["chunk_num"]
+        with open(Path(path), "r") as f:
+            raw_full = f.read()
+        chunks = splitter.split_text(raw_full)
+        text = chunks[chunk_num]
+        print(f"Similarity: {similarity.item():.4f}")
+        print(f"File:       {path}")
+        if int(start_line) == int(end_line):
+            print(f"Line:       {start_line}")
+        else:
+            print(f"Lines:      {start_line}-{end_line}")
+        print(f"Text:       {text}")
+        print("-------------------------------------------------------------")
+
+
+def _getQuery(embedder: embed_engine) -> str:
+    print("Query:")
+    raw = input("> ")
+    return embedder._embedQuery(raw)
+
+
+def _lookDir(
+    splitter: RecursiveCharacterTextSplitter,
+    path: Path,
+    dbM: dbManager,
+    embedder: embed_engine,
+    _seen=None,
+):
+    _seen = _seen if _seen is not None else set()
+    real = path.resolve()
+    if real in _seen:
+        return
+    _seen.add(real)
+    pathlist = []
+    for item in path.iterdir():
+        pathlist.append(item)
+        if item.is_dir():
+            _lookDir(splitter, item, dbM, embedder, _seen)
+        elif is_vectorizable(item):
+            _splitFile(item, splitter, dbM, embedder)
+
+
+def _listAll(
+    path: Path,
+    _seen=None,
+) -> list[Path]:
+    _seen = _seen if _seen is not None else set()
+    real = path.resolve()
+    if real in _seen:
+        return []
+    _seen.add(real)
+
+    paths = []
+    for item in path.iterdir():
+        if item.is_dir():
+            paths += _listAll(item, _seen)
+        elif is_vectorizable(item):
+            paths.append(item)
+    return paths
+
+
+def _initDB(connection: str) -> dbManager:
+    db = sqlite3.connect(connection)
+    db.execute("PRAGMA foreign_keys = ON")
+    return dbManager(db)
+
+
+def _splitFile(
+    file: Path,
+    splitter: RecursiveCharacterTextSplitter,
+    dbM: dbManager,
+    embedder: embed_engine,
+):
+    time = datetime.fromtimestamp(file.stat().st_mtime)
+    with open(file, "r") as f:
+        raw_full = f.read()
+    chunks = splitter.split_text(raw_full)
+    to_embed = []
+    entries = []
+    line = 0
+    for i, chunk in enumerate(chunks):
+        start_line = line
+        line += chunk.count("\n")
+        end_line = line
+
+        to_embed.append(chunk)
+        ent = entry(
+            path=file,
+            start_line=start_line,
+            end_line=end_line,
+            time=time,
+            chunk_num=i,
+        )
+        # ent.printAll()
+        entries.append(ent)
+    embedded = embedder.embedAll(to_embed)
+    for chunk, ent in zip(embedded, entries):
+        ent.embedding = chunk
+    dbM.insertFull(entries)
+    print(f"[DEBUG] inserted all chunks of {file.name}")