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

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: hie_rag
-Version: 0.1.3
+Version: 0.2.0
 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.0/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.0}/hie_rag/hie_rag.py

@@ -6,12 +6,12 @@ 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..."}

{hie_rag-0.1.3 → hie_rag-0.2.0}/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):
+        self.client = ChatOllama(model="llama3.2:latest")
+        self.utils = Utils(base_url=base_url)
 
     def _generate_metadata(self, chunk: str) -> Dict:
         """Generate metadata for a chunk using LangChain"""

{hie_rag-0.1.3 → hie_rag-0.2.0}/hie_rag/split.py

@@ -4,11 +4,11 @@ from .utils import Utils
 
 
 class Split:
-    def __init__(self, api_key: str = None):
+    def __init__(self, base_url: str = None):
         """
         Initializes the Split object with default or user-defined thresholds.
         """
-        self.utils = Utils(api_key=api_key)
+        self.utils = Utils(base_url=base_url)
 
     def _split_large_chunk(self, paragraphs: List[str], embeddings: List[List[float]]) -> (List[str], List[str]):
         """
@@ -34,8 +34,8 @@ class Split:
     def split(
         self,
         extracted_text: str,
-        min_chunk_size: int = 4000,
-        max_chunk_size: int = 7000
+        min_chunk_size: int = 300,
+        max_chunk_size: int = 500
     ) -> List[str]:
         """
         Splits the input text into chunks of token-size between [min_chunk_size, max_chunk_size].

{hie_rag-0.1.3 → hie_rag-0.2.0}/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.0}/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):
+        self.client = ChatOllama(model="llama3.2:latest")
+        self.utils = Utils(base_url=base_url)
 
     def _convert_to_string(self, chunk_metadata: dict) -> str:
         """

{hie_rag-0.1.3 → hie_rag-0.2.0}/hie_rag/utils.py

@@ -6,17 +6,20 @@ import tempfile
 import numpy as np
 import tiktoken
 from markitdown import MarkItDown
-from openai import OpenAI
 from sklearn.metrics.pairwise import cosine_similarity
 
+from .ai_client import AiClient
+
 
 class Utils:
-    def __init__(self, api_key=None):
-        self.client = OpenAI(api_key=api_key)
+    def __init__(self, base_url=None):
+        # 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(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):
@@ -46,18 +49,15 @@ class Utils:
         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="nomic-embed-text") -> list:
+        if not self.client:
+            raise RuntimeError("No embedding client configured")
+        return self.client.get_embedding(text, model=model)
 
-    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 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"""

{hie_rag-0.1.3 → hie_rag-0.2.0}/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.0}/hie_rag.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: hie_rag
-Version: 0.1.3
+Version: 0.2.0
 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.0}/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.0}/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.0}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "hie_rag"
-version = "0.1.3"
+version = "0.2.0"
 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.0}/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.0}/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.1.3 → hie_rag-0.2.0}/test/test-utils.py

@@ -1,11 +1,15 @@
 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(api_key=os.getenv("OPENAI_API_KEY"))
+utils = Utils(base_url=os.getenv("BASE_URL"))
 
 with open("test.pdf", "rb") as uploaded_file:
     extracted_text = utils.extract_text(uploaded_file)
@@ -27,7 +31,7 @@ result_get_embedding = utils.get_embedding(extracted_text[:100])
 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:
+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")

{hie_rag-0.1.3 → hie_rag-0.2.0}/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)}