hie-rag 0.1.3__tar.gz → 0.2.2__tar.gz

{hie_rag-0.1.3 → hie_rag-0.2.2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: hie_rag
-Version: 0.1.3
+Version: 0.2.2
 Summary: A hierarchical RAG framework for chunks retrieval.
 License: Apache License
         Version 2.0, January 2004
@@ -29,7 +29,7 @@ Requires-Dist: openai==1.66.3
 Requires-Dist: scikit-learn
 Requires-Dist: tiktoken==0.8.0
 Requires-Dist: langchain==0.3.13
-Requires-Dist: langchain-openai==0.2.13
+Requires-Dist: langchain-ollama==0.3.3
 Requires-Dist: chromadb==0.6.2
 Dynamic: license-file
 

hie_rag-0.2.2/hie_rag/ai_client.py

@@ -0,0 +1,26 @@
+import requests
+
+
+class AiClient:
+    def __init__(self, base_url="http://localhost:11434"):
+        self.base_url = base_url
+        self.headers = {"Content-Type": "application/json"}
+
+    def get_embedding(self, text: str, model="nomic-embed-text") -> list:
+        url = f"{self.base_url}/api/embeddings"
+        payload = {
+            "model": model,
+            "prompt": text
+        }
+        response = requests.post(url, json=payload, headers=self.headers, timeout=60)
+        response.raise_for_status()
+        data = response.json()
+
+        # Extract embedding, adapt if your API response structure differs
+        embedding = data.get("embedding") or (data.get("data") and data["data"][0].get("embedding"))
+        if embedding is None:
+            raise ValueError("Embedding not found in Ollama response")
+        return embedding
+
+    def list_embeddings(self, texts: list, model="nomic-embed-text") -> list:
+        return [self.get_embedding(text, model=model) for text in texts]

{hie_rag-0.1.3 → hie_rag-0.2.2}/hie_rag/hie_rag.py

@@ -6,27 +6,32 @@ from hie_rag.vectordb import Vectordb
 
 
 class HieRag:
-    def __init__(self, api_key=None, path="./db", collection_name="db_collection"):
-        self.split = Split(api_key=api_key)
-        self.utils = Utils(api_key=api_key)
-        self.tree_index = TreeIndex(api_key=api_key)
-        self.process = Process(api_key=api_key)
-        self.vector_db = Vectordb(path=path, api_key=api_key, collection_name=collection_name)
+    def __init__(self, base_url, path="./db", collection_name="db_collection"):
+        self.split = Split(base_url=base_url)
+        self.utils = Utils(base_url=base_url)
+        self.tree_index = TreeIndex(base_url=base_url)
+        self.process = Process(base_url=base_url)
+        self.vector_db = Vectordb(path=path, base_url=base_url, collection_name=collection_name)
     
     def process_and_save_index_stream(self, file_name: str, uploaded_file: bytes, min_chunk_size, max_chunk_size):
         yield {"status": "🔍 Extracting text..."}
+        print(f"Extracting text from {file_name}")
         extracted_text = self.utils.extract_text(uploaded_file)
 
         yield {"status": "✂️ Splitting into chunks..."}
+        print(f"Splitting text into chunks with min size {min_chunk_size} and max size {max_chunk_size}")
         result_split = self.split.split(extracted_text, min_chunk_size=min_chunk_size, max_chunk_size=max_chunk_size)
 
         yield {"status": "🧠 Processing chunks..."}
+        print(f"Processing {len(result_split)} chunks")
         result_process = self.process.process_chunks(result_split)
 
         yield {"status": "🌲 Building tree index..."}
+        print(f"Building tree index with {len(result_process)} chunks")
         tree_index = self.tree_index.tree_index(file_name = file_name, chunk_metadata=result_process)
 
         yield {"status": "💾 Saving to vector DB..."}
+        print(f"Saving tree index with {len(tree_index.get('chunks', []))} chunks to vector DB")
         save_result = self.vector_db.save_index(tree_index)
 
         file_id = save_result.get("file_id", "unknown")

{hie_rag-0.1.3 → hie_rag-0.2.2}/hie_rag/process.py

@@ -1,7 +1,7 @@
 from typing import Dict, List
 
 from langchain_core.prompts import PromptTemplate
-from langchain_openai import ChatOpenAI
+from langchain_ollama import ChatOllama
 from pydantic import Field
 from typing_extensions import TypedDict
 
@@ -9,9 +9,9 @@ from .utils import Utils
 
 
 class Process:
-    def __init__(self, api_key=None):
-        self.client = ChatOpenAI(temperature=0, model="gpt-4o", api_key=api_key)
-        self.utils = Utils(api_key=api_key)
+    def __init__(self, base_url=None, model="llama3.2:latest"):
+        self.client = ChatOllama(model=model)
+        self.utils = Utils(base_url=base_url)
 
     def _generate_metadata(self, chunk: str) -> Dict:
         """Generate metadata for a chunk using LangChain"""

hie_rag-0.2.2/hie_rag/split.py

@@ -0,0 +1,129 @@
+import re
+from collections import deque
+from typing import List, Tuple
+
+from .utils import Utils
+
+
+class Split:
+    def __init__(self, base_url: str = None):
+        """
+        Initializes the Split object with default or user-defined thresholds.
+        """
+        self.utils = Utils(base_url=base_url)
+
+    def _custom_split(self, text: str):
+        stripped = text.strip()
+        # 以「空白行」作為段落切點
+        raw_paragraphs = re.split(r'\n\s*\n+', stripped)
+
+        result = []
+        for para in raw_paragraphs:
+            # 把段落內所有換行改成空格
+            single_line = para.replace('\r\n', ' ').replace('\r', ' ').replace('\n', ' ')
+            cleaned = single_line.strip()
+            if cleaned:
+                result.append(cleaned)
+        return result
+
+    def _split_large_chunk(self, paragraphs: List[str], embeddings: List[List[float]]) -> (List[str], List[str]):
+        """
+        Splits 'paragraphs' by finding the least similar boundary using 'embeddings'
+        (which are precomputed for these paragraphs only). Returns (left_part, right_part).
+        """
+        # If there are 0 or 1 paragraphs, no need to split
+        if len(paragraphs) < 2:
+            return paragraphs, []
+
+        # We'll assume 'embeddings' is already the same length as 'paragraphs'.
+        if len(embeddings) < 2:
+            # Can't compute consecutive similarities with fewer than 2 embeddings
+            return paragraphs, []
+
+        # Find the least similar consecutive boundary
+        window_size = 3
+        split_index = self.utils.get_windowed_least_similar(embeddings, window_size=window_size)
+        
+        left_part = paragraphs[:split_index + 1]
+        right_part = paragraphs[split_index + 1:]
+        return left_part, right_part
+
+    def split(
+        self,
+        extracted_text: str,
+        min_chunk_size: int = 300,
+        max_chunk_size: int = 500
+    ) -> List[str]:
+        
+        # 1) Build a deque of triples, so we never mutate three separate lists:
+        # paras = [p.strip() for p in extracted_text.split("\n\n") if p.strip()]
+        paras = self._custom_split(extracted_text)
+        
+        if not paras:
+            return []
+
+        tokens = [self.utils.count_tokens(p) for p in paras]
+        embs   = self.utils.list_embeddings(paras)
+        D: deque[Tuple[str,List[float],int]] = deque(
+            zip(paras, embs, tokens)
+        )
+
+        final_chunks: List[str] = []
+
+        # 2) As long as there’s anything left in D, build one chunk at a time:
+        while D:
+            cur_paras:    List[str]           = []
+            cur_embs:     List[List[float]]   = []
+            cur_tokens:   List[int]           = []
+            total_tokens = 0
+
+            # 2a) Guarantee we hit at least min_chunk_size
+            while D and total_tokens < min_chunk_size:
+                p, e, t = D.popleft()
+                # if even this one p would bust max, you might choose to take it alone
+                if total_tokens + t > max_chunk_size and total_tokens > 0:
+                    # push it back for the next round
+                    D.appendleft((p,e,t))
+                    break
+                cur_paras.append(p)
+                cur_embs .append(e)
+                cur_tokens.append(t)
+                total_tokens += t
+
+            # if we ran out before min and have something -> emit it
+            if total_tokens < min_chunk_size and not D:
+                final_chunks.append(" ".join(cur_paras))
+                break
+
+            # 2b) Greedily fill until just under max_chunk_size
+            while D and total_tokens + D[0][2] <= max_chunk_size:
+                p, e, t = D.popleft()
+                cur_paras.append(p)
+                cur_embs .append(e)
+                cur_tokens.append(t)
+                total_tokens += t
+
+            # 3) Now we have between min and max tokens: split at the least-similar boundary
+            if cur_paras:
+                left, right = self._split_large_chunk(cur_paras, cur_embs)
+
+                # Count tokens in “left” to see if it meets min_chunk_size
+                left_token_count = sum(self.utils.count_tokens(p) for p in left)
+
+                if left_token_count >= min_chunk_size:
+                    # If left is big enough, emit it
+                    final_chunks.append(" ".join(left))
+
+                    # Push “right” (the remainder) back onto D for subsequent chunks
+                    for rp, re, rt in reversed(list(zip(
+                                    cur_paras[len(left):],
+                                    cur_embs  [len(left):],
+                                    cur_tokens[len(left):]
+                                ))):
+                        D.appendleft((rp, re, rt))
+                else:
+                    # If “left” is too small, just emit the entire cur_paras as one chunk
+                    final_chunks.append(" ".join(cur_paras))
+                    # (We do NOT push anything back, because cur_paras is fully consumed.)
+
+        return final_chunks 

{hie_rag-0.1.3 → hie_rag-0.2.2}/hie_rag/split_and_process.py

@@ -4,10 +4,10 @@ from hie_rag.utils import Utils
 
 
 class SplitAndProcess:
-    def __init__(self, api_key=None):
-        self.split = Split(api_key=api_key)
-        self.utils = Utils(api_key=api_key)
-        self.process = Process(api_key=api_key)
+    def __init__(self, base_url: str):
+        self.split = Split(base_url=base_url)
+        self.utils = Utils(base_url=base_url)
+        self.process = Process(base_url=base_url)
 
     def split_and_process(self, uploaded_file):
         extracted_text = self.utils.extract_text(uploaded_file)

{hie_rag-0.1.3 → hie_rag-0.2.2}/hie_rag/tree_index.py

@@ -2,7 +2,7 @@ import json
 from typing import List
 
 from langchain_core.prompts import PromptTemplate
-from langchain_openai import ChatOpenAI
+from langchain_ollama import ChatOllama
 from pydantic import Field
 from typing_extensions import TypedDict
 
@@ -10,9 +10,9 @@ from .utils import Utils
 
 
 class TreeIndex:
-    def __init__(self, api_key: str):
-        self.client = ChatOpenAI(temperature=0, model="gpt-4o", api_key=api_key)
-        self.utils = Utils(api_key=api_key)
+    def __init__(self, base_url: str, model="llama3.2:latest"):
+        self.client = ChatOllama(model=model)
+        self.utils = Utils(base_url=base_url)
 
     def _convert_to_string(self, chunk_metadata: dict) -> str:
         """
@@ -32,8 +32,8 @@ class TreeIndex:
 
             NOTE: 
             1. 請輸出繁體中文
-            2. The summary should be concise and capture the main points of the text.
-            3. The summary should be around 5-8 sentences long.
+            2. The summary should be concise with details and better than the individual summaries.
+            3. The summary should be long enough to cover all the main points of the text.
 
             Summaries:
             {summaries}

hie_rag-0.2.2/hie_rag/utils.py

@@ -0,0 +1,122 @@
+import contextlib
+import io
+import os
+import tempfile
+
+import numpy as np
+import tiktoken
+from markitdown import MarkItDown
+from sklearn.metrics.pairwise import cosine_similarity
+
+from .ai_client import AiClient
+
+
+class Utils:
+    def __init__(self, base_url: str):
+        # self.client = OpenAI(api_key=api_key)
+        self.client = AiClient(base_url=base_url)
+
+    def extract_text(self, uploaded_file: bytes):
+        """Extract text from an uploaded file using MarkItDown."""
+        # md = MarkItDown(llm_client=self.client, llm_model="gpt-4o")
+        md = MarkItDown()
+
+        # Accept both raw bytes and file-like objects with `.read()`
+        if isinstance(uploaded_file, bytes):
+            file_bytes = uploaded_file
+            suffix = ".bin"  # fallback generic extension
+        elif hasattr(uploaded_file, "read"):
+            file_bytes = uploaded_file.read()
+            filename = getattr(uploaded_file, "name", None) or getattr(uploaded_file, "filename", None)
+            suffix = os.path.splitext(filename)[-1] if filename else ".bin"
+        else:
+            raise TypeError("Unsupported file type: must be bytes or file-like object")
+
+        # Write to temp file for MarkItDown to process
+        with tempfile.NamedTemporaryFile(delete=False, suffix=suffix) as temp_file:
+            temp_file_path = temp_file.name
+            temp_file.write(file_bytes)
+
+        try:
+            # Redirect stderr to suppress native print warnings like "CropBox missing"
+            with contextlib.redirect_stderr(io.StringIO()):
+                extracted_text = md.convert(temp_file_path)
+        finally:
+            # Clean up the temporary file
+            os.remove(temp_file_path)
+
+        return extracted_text.text_content
+
+    def count_tokens(self, text: str, encoding="cl100k_base") -> int:
+        """Count tokens in text using tiktoken"""
+        tokenizer = tiktoken.get_encoding(encoding)
+        return len(tokenizer.encode(text))
+
+    def get_embedding(self, text: str, model="nomic-embed-text") -> list:
+        if not self.client:
+            raise RuntimeError("No embedding client configured")
+        return self.client.get_embedding(text, model=model)
+
+    def list_embeddings(self, chunks: list, model="nomic-embed-text") -> list:
+        if not self.client:
+            raise RuntimeError("No embedding client configured")
+        return self.client.list_embeddings(chunks, model=model)
+
+    def get_consecutive_least_similar(self, embeddings: list) -> int:
+        """Find the index where consecutive similarity is lowest"""
+        cs = cosine_similarity(embeddings)
+        
+        # Get similarities between consecutive sentences only
+        consecutive_similarities = []
+        for i in range(len(cs) - 1):
+            consecutive_similarities.append(cs[i][i + 1])
+        
+        # Find the index where consecutive similarity is lowest
+        split_index = np.argmin(consecutive_similarities)
+        
+        return split_index
+    
+    def get_windowed_least_similar(
+        self,
+        embeddings: list,
+        window_size: int = 3
+    ) -> int:
+        """
+        對 embeddings 做滑動窗口：對每個可能的分割位置 i（0 <= i < len-1），
+        將 [max(0, i-window_size+1) .. i] 這 window_size 句平均後的向量
+        與 [i+1 .. min(len, i+window_size)] 這 window_size 句平均後的向量做 cosine 相似度，
+        回傳相似度最低的那個 i。
+        """
+        if len(embeddings) < 2:
+            # 根本沒得分割
+            return 0
+
+        # 把 list-of-lists 轉成 numpy array (shape: [n_sentences, dim_emb])
+        embs = np.array(embeddings)
+        n = embs.shape[0]
+
+        best_index = 0
+        lowest_sim = float('inf')
+
+        for i in range(n - 1):
+            # 前半段：從 pre_start 到 i (inclusive)
+            pre_start = max(0, i - window_size + 1)
+            pre_group = embs[pre_start : i + 1]   # shape: (<=window_size, dim)
+
+            # 後半段：從 i+1 到 post_end-1
+            post_end = min(n, i + 1 + window_size)
+            post_group = embs[i + 1 : post_end]   # shape: (<=window_size, dim)
+
+            # 計算平均向量
+            # （也可以改成加總：np.sum(...)；不過平均比較常見且 scale 感覺一致）
+            pre_avg = np.mean(pre_group, axis=0).reshape(1, -1)   # shape: (1, dim)
+            post_avg = np.mean(post_group, axis=0).reshape(1, -1) # shape: (1, dim)
+
+            # 計算 cosine similarity
+            sim = float(cosine_similarity(pre_avg, post_avg)[0][0])
+
+            if sim < lowest_sim:
+                lowest_sim = sim
+                best_index = i
+
+        return best_index

{hie_rag-0.1.3 → hie_rag-0.2.2}/hie_rag/vectordb.py

@@ -7,9 +7,9 @@ from .utils import Utils
 
 
 class Vectordb():
-    def __init__(self, path, api_key, collection_name):
+    def __init__(self, path, base_url, collection_name):
         self.client = chromadb.PersistentClient(path = path)
-        self.utils = Utils(api_key=api_key)
+        self.utils = Utils(base_url=base_url)
         self.collection = self.client.get_or_create_collection(collection_name)
 
     def _convert_numpy(self, obj):

{hie_rag-0.1.3 → hie_rag-0.2.2}/hie_rag.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: hie_rag
-Version: 0.1.3
+Version: 0.2.2
 Summary: A hierarchical RAG framework for chunks retrieval.
 License: Apache License
         Version 2.0, January 2004
@@ -29,7 +29,7 @@ Requires-Dist: openai==1.66.3
 Requires-Dist: scikit-learn
 Requires-Dist: tiktoken==0.8.0
 Requires-Dist: langchain==0.3.13
-Requires-Dist: langchain-openai==0.2.13
+Requires-Dist: langchain-ollama==0.3.3
 Requires-Dist: chromadb==0.6.2
 Dynamic: license-file
 

{hie_rag-0.1.3 → hie_rag-0.2.2}/hie_rag.egg-info/SOURCES.txt

@@ -2,7 +2,7 @@ LICENSE
 README.md
 pyproject.toml
 hie_rag/__init__.py
-hie_rag/app.py
+hie_rag/ai_client.py
 hie_rag/hie_rag.py
 hie_rag/process.py
 hie_rag/split.py

{hie_rag-0.1.3 → hie_rag-0.2.2}/hie_rag.egg-info/requires.txt

@@ -6,5 +6,5 @@ openai==1.66.3
 scikit-learn
 tiktoken==0.8.0
 langchain==0.3.13
-langchain-openai==0.2.13
+langchain-ollama==0.3.3
 chromadb==0.6.2

{hie_rag-0.1.3 → hie_rag-0.2.2}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "hie_rag"
-version = "0.1.3"
+version = "0.2.2"
 description = "A hierarchical RAG framework for chunks retrieval."
 license = { file = "LICENSE" }
 dependencies = [
@@ -16,7 +16,7 @@ dependencies = [
     "scikit-learn",
     "tiktoken==0.8.0",
     "langchain==0.3.13",
-    "langchain-openai==0.2.13",
+    "langchain-ollama==0.3.3",
     "chromadb==0.6.2"
 ]
 readme = "README.md"

{hie_rag-0.1.3 → hie_rag-0.2.2}/test/test-process.py

@@ -1,15 +1,19 @@
 import os
+import sys
+
+sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))
 
 from dotenv import load_dotenv
-from hie_rag.app import Split
+
 from hie_rag.process import Process
+from hie_rag.split import Split
 from hie_rag.utils import Utils
 
 load_dotenv()
 
-split = Split(api_key=os.getenv("OPENAI_API_KEY"), min_chunk_size=200, max_chunk_size=500)
-utils = Utils(api_key=os.getenv("OPENAI_API_KEY"))
-process = Process(api_key=os.getenv("OPENAI_API_KEY"))
+split = Split(base_url=os.getenv('BASE_URL'))
+utils = Utils(base_url=os.getenv('BASE_URL'))
+process = Process(base_url=os.getenv('BASE_URL'))
 
 with open("test.pdf", "rb") as uploaded_file:
     extracted_text = utils.extract_text(uploaded_file)
@@ -20,7 +24,7 @@ result_split = split.split(extracted_text)
 result_process = process.process_chunks(result_split)
 
 # Write results to the text file
-with open("test-process-result", "w", encoding="utf-8") as file:
+with open("test-process-result-new", "w", encoding="utf-8") as file:
     file.write("Processed Chunks:\n")
     file.write(str(result_process) + "\n")
 

{hie_rag-0.1.3 → hie_rag-0.2.2}/test/test-split.py

@@ -1,13 +1,17 @@
 import os
+import sys
+
+sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))
 
 from dotenv import load_dotenv
-from hie_rag.app import Split
+
+from hie_rag.split import Split
 from hie_rag.utils import Utils
 
 load_dotenv()
 
-split = Split(api_key=os.getenv("OPENAI_API_KEY"), min_chunk_size=200, max_chunk_size=500)
-utils = Utils(api_key=os.getenv("OPENAI_API_KEY"))
+split = Split(base_url=os.getenv("BASE_URL"))
+utils = Utils(base_url=os.getenv("BASE_URL"))
 
 with open("test.pdf", "rb") as uploaded_file:
     extracted_text = utils.extract_text(uploaded_file)
@@ -16,10 +20,10 @@ with open("test.pdf", "rb") as uploaded_file:
 extracted_text = extracted_text[:1000]
 
 # Split the extracted text
-result_split = split.split(extracted_text)
+result_split = split.split(extracted_text, min_chunk_size=300, max_chunk_size=500)
 
 # Write results to the text file
-with open("test-split-result", "w", encoding="utf-8") as file:
+with open("test-split-result-new", "w", encoding="utf-8") as file:
     file.write("Splitted Text:\n")
     file.write(str(result_split) + "\n")
     file.write("Length of the Splitted Text:\n")

hie_rag-0.2.2/test/test-utils.py

@@ -0,0 +1,56 @@
+import os
+import sys
+
+sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))
+
+from dotenv import load_dotenv
+
+from hie_rag.utils import Utils
+
+load_dotenv()
+
+utils = Utils(base_url=os.getenv("BASE_URL"))
+
+with open("test2.pdf", "rb") as uploaded_file:
+    extracted_text = utils.extract_text(uploaded_file)
+
+# Count tokens for the first 100 words
+# result_count_tokens = utils.count_tokens(extracted_text[:100])
+result_count_tokens = utils.count_tokens(extracted_text)
+print(f"Token count: {result_count_tokens}")
+
+# # Get embeddings for the text slices
+# result_list_embeddings = utils.list_embeddings([
+#     extracted_text[:100],
+#     extracted_text[100:200],
+#     extracted_text[200:300]
+# ])
+
+# # Get the embedding for the first 100 words
+# result_get_embedding = utils.get_embedding(extracted_text[:100])
+
+# # Find the index of least similar consecutive embeddings
+# result_get_consecutive_least_similar = utils.get_consecutive_least_similar(result_list_embeddings)
+
+# # Write results to the text file
+# with open("test-utils-result-new", "w", encoding="utf-8") as file:
+#     file.write("Extracted Text:\n")
+#     file.write(extracted_text + "\n\n")
+#     file.write("====================================\n\n")
+    
+#     file.write("Count of Tokens (First 100 words):\n")
+#     file.write(str(result_count_tokens) + "\n\n")
+#     file.write("====================================\n\n")
+    
+#     file.write("List of Embeddings:\n")
+#     file.write(str(result_list_embeddings) + "\n\n")
+#     file.write("====================================\n\n")
+    
+#     file.write("Embedding of First 100 words:\n")
+#     file.write(str(result_get_embedding) + "\n\n")
+#     file.write("====================================\n\n")
+    
+#     file.write("Index of Least Similar Consecutive Embeddings:\n")
+#     file.write(str(result_get_consecutive_least_similar) + "\n")
+
+# print("Results written to a txt file.")

{hie_rag-0.1.3 → hie_rag-0.2.2}/test/test_split_and_process.py

@@ -1,16 +1,21 @@
 import os
+import sys
+
+sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))
+
 
 from dotenv import load_dotenv
+
 from hie_rag import SplitAndProcess
 
 load_dotenv()
 
-split_and_process = SplitAndProcess(api_key=os.getenv("OPENAI_API_KEY"))
+split_and_process = SplitAndProcess(base_url=os.getenv("BASE_URL"))
 
 with open("test.pdf", "rb") as uploaded_file:
     result_process = split_and_process.split_and_process(uploaded_file)
 
-with open("test-split-and-process-result", "w", encoding="utf-8") as file:
+with open("test-split-and-process-result-new", "w", encoding="utf-8") as file:
     file.write("Split and Processed Text:\n")
     file.write(str(result_process) + "\n")
 

hie_rag-0.1.3/hie_rag/app.py

@@ -1,77 +0,0 @@
-# import json
-# import os
-
-# from .generate import Generate
-# from .process import Process
-# from .split import Split
-# from .tree_index import TreeIndex
-# from .utils import Utils
-# from .vectordb import Vectordb
-
-
-# # Function to handle data
-# def handle_data(data):
-#     """
-#     Processes incoming data and returns a response.
-#     """
-#     try:
-#         # This is the logic that used to be in the /api/data route
-#         return {"received": data}
-#     except Exception as e:
-#         return {"error": str(e)}
-
-# # Function to handle file upload and processing
-# def handle_file_upload(uploaded_file, access_token):
-#     """
-#     Processes the uploaded file and extracts its text.
-#     """
-#     try:
-#         utils = Utils(api_key=access_token)
-#         process = Process(api_key=access_token)
-#         split = Split(api_key=access_token)
-#         tree_index = TreeIndex(api_key=access_token)
-
-#         if uploaded_file is None:
-#             return {"error": "No file selected for uploading"}
-
-#         filename = uploaded_file.filename
-#         extracted_text = utils.extract_text(uploaded_file)
-#         final_chunk_list = split.split(extracted_text)
-#         processed_chunks = process.process_chunks(final_chunk_list)
-#         data = tree_index.output_index(processed_chunks)
-
-#         return {"filename": filename, "data": data}
-#     except Exception as e:
-#         return {"error": str(e)}
-
-# # Function to handle generation logic
-# def handle_generation(file, access_token):
-#     """
-#     Handles the file for generation and returns generated data.
-#     """
-#     try:
-#         data = json.load(file)
-
-#         if "chunks" not in data:
-#             return {"error": "Missing 'chunks' in data"}
-
-#         path = os.getenv("INDEX_PATH")
-#         vectordb = Vectordb(path=path, api_key=access_token)
-#         generate = Generate(api_key=access_token)
-
-#         save_index_result = vectordb.save_index(data)
-#         generated_full_data = []
-        
-#         for i in data["chunks"]:
-#             original_chunk = i["original_chunk"]
-#             query_result = vectordb.query_by_text(original_chunk, n_results=3)
-#             possible_reference = query_result["metadatas"][0][1]["summary"] + "\n" + query_result["metadatas"][0][2]["summary"]
-
-#             data_gen = generate.generate(original_chunk, possible_reference)
-#             generated_full_data.extend(data_gen["dataset"])
-
-#         return {"data": generated_full_data}
-#     except json.JSONDecodeError:
-#         return {"error": "Invalid JSON file format"}
-#     except Exception as e:
-#         return {"error": str(e)}

hie_rag-0.1.3/hie_rag/split.py

@@ -1,120 +0,0 @@
-from typing import List
-
-from .utils import Utils
-
-
-class Split:
-    def __init__(self, api_key: str = None):
-        """
-        Initializes the Split object with default or user-defined thresholds.
-        """
-        self.utils = Utils(api_key=api_key)
-
-    def _split_large_chunk(self, paragraphs: List[str], embeddings: List[List[float]]) -> (List[str], List[str]):
-        """
-        Splits 'paragraphs' by finding the least similar boundary using 'embeddings'
-        (which are precomputed for these paragraphs only). Returns (left_part, right_part).
-        """
-        # If there are 0 or 1 paragraphs, no need to split
-        if len(paragraphs) < 2:
-            return paragraphs, []
-
-        # We'll assume 'embeddings' is already the same length as 'paragraphs'.
-        if len(embeddings) < 2:
-            # Can't compute consecutive similarities with fewer than 2 embeddings
-            return paragraphs, []
-
-        # Find the least similar consecutive boundary
-        split_index = self.utils.get_consecutive_least_similar(embeddings)
-
-        left_part = paragraphs[:split_index + 1]
-        right_part = paragraphs[split_index + 1:]
-        return left_part, right_part
-
-    def split(
-        self,
-        extracted_text: str,
-        min_chunk_size: int = 4000,
-        max_chunk_size: int = 7000
-    ) -> List[str]:
-        """
-        Splits the input text into chunks of token-size between [min_chunk_size, max_chunk_size].
-        Once a chunk is in that range, we find the "least similar" boundary, store the left side,
-        and re-insert the right side for further splitting.
-        """
-        paragraphs = [p.strip() for p in extracted_text.split("\n\n") if p.strip()]
-        if not paragraphs:
-            return []
-
-        # Precompute once
-        paragraphs_tokens = [self.utils.count_tokens(p) for p in paragraphs]
-        paragraphs_embeddings = self.utils.list_embeddings(paragraphs)
-
-        final_chunks = []
-        idx = 0
-        n = len(paragraphs)
-
-        while idx < n:
-            chunk_paragraphs = []
-            chunk_embeddings = []
-            chunk_tokens = []  # Keep track of tokens in this chunk
-            current_tokens = 0
-
-            # 1) Accumulate until we at least exceed min_chunk_size or run out
-            while idx < n and current_tokens < min_chunk_size:
-                if current_tokens + paragraphs_tokens[idx] <= max_chunk_size:
-                    chunk_paragraphs.append(paragraphs[idx])
-                    chunk_embeddings.append(paragraphs_embeddings[idx])
-                    chunk_tokens.append(paragraphs_tokens[idx])
-                    current_tokens += paragraphs_tokens[idx]
-                    idx += 1
-                else:
-                    # This paragraph alone might exceed max_chunk_size -> handle as you see fit
-                    break
-
-            # If we haven't hit min_chunk_size but are out of paragraphs, store remainder and quit
-            if current_tokens < min_chunk_size and idx >= n:
-                if chunk_paragraphs:
-                    final_chunks.append(" ".join(chunk_paragraphs))
-                break
-
-            # 2) Keep adding while we're under max_chunk_size
-            while idx < n:
-                if current_tokens + paragraphs_tokens[idx] <= max_chunk_size:
-                    chunk_paragraphs.append(paragraphs[idx])
-                    chunk_embeddings.append(paragraphs_embeddings[idx])
-                    chunk_tokens.append(paragraphs_tokens[idx])
-                    current_tokens += paragraphs_tokens[idx]
-                    idx += 1
-                else:
-                    break
-
-            # Now we have between min_chunk_size and max_chunk_size tokens in 'chunk_paragraphs'
-            if chunk_paragraphs:
-                # 3) Split at the "least similar" boundary
-                left_part, right_part = self._split_large_chunk(
-                    chunk_paragraphs, chunk_embeddings
-                )
-
-                # We'll figure out how many paragraphs ended up in the left part
-                used_count = len(left_part)
-                leftover_count = len(right_part)
-
-                # Store left side
-                final_chunks.append(" ".join(left_part))
-
-                # If there's leftover, reinsert it into the main lists
-                if leftover_count > 0:
-                    # Slices for leftover
-                    leftover_embeddings = chunk_embeddings[used_count:]
-                    leftover_tokens = chunk_tokens[used_count:]
-
-                    # Re-insert them at index=idx
-                    paragraphs[idx:idx] = right_part
-                    paragraphs_embeddings[idx:idx] = leftover_embeddings
-                    paragraphs_tokens[idx:idx] = leftover_tokens
-
-                    # Recompute n, in case the paragraphs list has grown
-                    n = len(paragraphs)
-
-        return final_chunks

hie_rag-0.1.3/hie_rag/utils.py

@@ -1,74 +0,0 @@
-import contextlib
-import io
-import os
-import tempfile
-
-import numpy as np
-import tiktoken
-from markitdown import MarkItDown
-from openai import OpenAI
-from sklearn.metrics.pairwise import cosine_similarity
-
-
-class Utils:
-    def __init__(self, api_key=None):
-        self.client = OpenAI(api_key=api_key)
-
-    def extract_text(self, uploaded_file: bytes):
-        """Extract text from an uploaded file using MarkItDown."""
-        md = MarkItDown(llm_client=self.client, llm_model="gpt-4o")
-
-        # Accept both raw bytes and file-like objects with `.read()`
-        if isinstance(uploaded_file, bytes):
-            file_bytes = uploaded_file
-        elif hasattr(uploaded_file, "read"):
-            file_bytes = uploaded_file.read()
-        else:
-            raise TypeError("Unsupported file type: must be bytes or file-like object")
-
-        # Write to temp file for MarkItDown to process
-        with tempfile.NamedTemporaryFile(delete=False) as temp_file:
-            temp_file_path = temp_file.name
-            temp_file.write(file_bytes)
-
-        try:
-            # Redirect stderr to suppress native print warnings like "CropBox missing"
-            with contextlib.redirect_stderr(io.StringIO()):
-                extracted_text = md.convert(temp_file_path)
-        finally:
-            # Clean up the temporary file
-            os.remove(temp_file_path)
-
-        return extracted_text.text_content
-
-    def count_tokens(self, text: str, encoding="cl100k_base") -> int:
-        """Count tokens in text using tiktoken"""
-        tokenizer = tiktoken.get_encoding(encoding)
-        return len(tokenizer.encode(text))
-
-    def list_embeddings(self, chunks: list, model="text-embedding-3-small") -> list:
-        """Get embeddings for a list of text chunks"""
-        embeddings = []
-        for chunk in chunks:
-            response = self.client.embeddings.create(input=chunk, model=model)
-            embeddings.append(response.data[0].embedding)
-        return embeddings
-
-    def get_embedding(self, text: str, model="text-embedding-3-small") -> list:
-        """Get embedding for a text"""
-        response = self.client.embeddings.create(input=text, model=model)
-        return response.data[0].embedding
-
-    def get_consecutive_least_similar(self, embeddings: list) -> int:
-        """Find the index where consecutive similarity is lowest"""
-        cs = cosine_similarity(embeddings)
-        
-        # Get similarities between consecutive sentences only
-        consecutive_similarities = []
-        for i in range(len(cs) - 1):
-            consecutive_similarities.append(cs[i][i + 1])
-        
-        # Find the index where consecutive similarity is lowest
-        split_index = np.argmin(consecutive_similarities)
-        
-        return split_index

hie_rag-0.1.3/test/test-utils.py

@@ -1,50 +0,0 @@
-import os
-
-from dotenv import load_dotenv
-from hie_rag.utils import Utils
-
-load_dotenv()
-
-utils = Utils(api_key=os.getenv("OPENAI_API_KEY"))
-
-with open("test.pdf", "rb") as uploaded_file:
-    extracted_text = utils.extract_text(uploaded_file)
-
-# Count tokens for the first 100 words
-result_count_tokens = utils.count_tokens(extracted_text[:100])
-
-# Get embeddings for the text slices
-result_list_embeddings = utils.list_embeddings([
-    extracted_text[:100],
-    extracted_text[100:200],
-    extracted_text[200:300]
-])
-
-# Get the embedding for the first 100 words
-result_get_embedding = utils.get_embedding(extracted_text[:100])
-
-# Find the index of least similar consecutive embeddings
-result_get_consecutive_least_similar = utils.get_consecutive_least_similar(result_list_embeddings)
-
-# Write results to the text file
-with open("test-utils-result", "w", encoding="utf-8") as file:
-    file.write("Extracted Text:\n")
-    file.write(extracted_text + "\n\n")
-    file.write("====================================\n\n")
-    
-    file.write("Count of Tokens (First 100 words):\n")
-    file.write(str(result_count_tokens) + "\n\n")
-    file.write("====================================\n\n")
-    
-    file.write("List of Embeddings:\n")
-    file.write(str(result_list_embeddings) + "\n\n")
-    file.write("====================================\n\n")
-    
-    file.write("Embedding of First 100 words:\n")
-    file.write(str(result_get_embedding) + "\n\n")
-    file.write("====================================\n\n")
-    
-    file.write("Index of Least Similar Consecutive Embeddings:\n")
-    file.write(str(result_get_consecutive_least_similar) + "\n")
-
-print("Results written to a txt file.")