ragxo 0.1.0__tar.gz

ragxo-0.1.0/PKG-INFO

@@ -0,0 +1,30 @@
+Metadata-Version: 2.1
+Name: ragxo
+Version: 0.1.0
+Summary: A RAG (Retrieval-Augmented Generation) toolkit with Milvus integration
+Home-page: https://github.com/yourusername/ragx
+License: MIT
+Keywords: rag,milvus,nlp,embeddings,openai
+Author: Mohamed Sadek
+Author-email: mohamedfawzydes@gmail.com
+Requires-Python: >=3.11,<4.0
+Classifier: Development Status :: 3 - Alpha
+Classifier: Intended Audience :: Developers
+Classifier: License :: OSI Approved :: MIT License
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Classifier: Programming Language :: Python :: 3.13
+Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
+Classifier: Topic :: Software Development :: Libraries :: Python Modules
+Requires-Dist: dill (>=0.3.9,<0.4.0)
+Requires-Dist: milvus (>=2.3.9,<3.0.0)
+Requires-Dist: nltk (>=3.9.1,<4.0.0)
+Requires-Dist: openai (>=1.61.1,<2.0.0)
+Requires-Dist: pydantic (>=2.10.6,<3.0.0)
+Requires-Dist: pymilvus (>=2.5.4,<3.0.0)
+Requires-Dist: sentence-transformers (>=3.4.1,<4.0.0)
+Project-URL: Repository, https://github.com/yourusername/ragx
+Description-Content-Type: text/markdown
+
+

ragxo-0.1.0/pyproject.toml

@@ -0,0 +1,53 @@
+[tool.poetry]
+name = "ragxo"
+version = "0.1.0"
+description = "A RAG (Retrieval-Augmented Generation) toolkit with Milvus integration"
+authors = ["Mohamed Sadek <mohamedfawzydes@gmail.com>"]
+readme = "README.md"
+license = "MIT"
+repository = "https://github.com/yourusername/ragx"
+keywords = ["rag", "milvus", "nlp", "embeddings", "openai"]
+classifiers = [
+    "Development Status :: 3 - Alpha",
+    "Intended Audience :: Developers",
+    "Topic :: Software Development :: Libraries :: Python Modules",
+    "Topic :: Scientific/Engineering :: Artificial Intelligence",
+    "License :: OSI Approved :: MIT License",
+    "Programming Language :: Python :: 3",
+    "Programming Language :: Python :: 3.11",
+]
+packages = [{include = "ragxo"}]
+
+[tool.poetry.dependencies]
+python = "^3.11"
+milvus = "^2.3.9"
+pymilvus = "^2.5.4"
+pydantic = "^2.10.6"
+nltk = "^3.9.1"
+sentence-transformers = "^3.4.1"
+dill = "^0.3.9"
+openai = "^1.61.1"
+
+[tool.poetry.group.dev.dependencies]
+pytest = "^8.3.4"
+black = "^24.2.0"
+isort = "^5.13.2"
+flake8 = "^7.0.0"
+mypy = "^1.9.0"
+
+[build-system]
+requires = ["poetry-core"]
+build-backend = "poetry.core.masonry.api"
+
+[tool.black]
+line-length = 88
+target-version = ['py311']
+
+[tool.isort]
+profile = "black"
+multi_line_output = 3
+
+[tool.mypy]
+python_version = "3.11"
+strict = true
+ignore_missing_imports = true

ragxo-0.1.0/ragxo/ragx.py

@@ -0,0 +1,141 @@
+from typing import Self, Callable
+from pymilvus import MilvusClient
+from pydantic import BaseModel
+import dill
+import os
+import shutil
+import logging
+import openai
+from openai import ChatCompletion
+
+logging.basicConfig(level=logging.DEBUG)
+logger = logging.getLogger(__name__)
+
+class Document(BaseModel):
+    text: str
+    metadata: dict
+    id: int
+
+class Ragx:
+    def __init__(self, dimension: int) -> None:
+        self.dimension = dimension
+        self.collection_name = "ragx"
+        self.db_path = "milvus.db"
+        self.client = MilvusClient(self.db_path)
+        self.client.create_collection(self.collection_name, dimension=dimension)
+        self.processing_fn = []
+        self.embedding_fn = None
+        self.system_prompt = None
+        self.model = "gpt-4o-mini"
+    
+    def add_preprocess(self, fn: Callable) -> Self:
+        self.processing_fn.append(fn)
+        return self
+    
+    def add_embedding_fn(self, fn: Callable) -> Self:
+        if not fn:
+            raise ValueError("Embedding function cannot be None")
+        self.embedding_fn = fn
+        return self
+    
+    def add_system_prompt(self, prompt: str) -> Self:
+        self.system_prompt = prompt
+        return self
+    
+    def add_model(self, model: str) -> Self:
+        self.model = model
+        return self
+    
+    def index(self, data: list[Document]) -> Self:
+        if not self.embedding_fn:
+            raise ValueError("Embedding function not set")
+            
+        processed_text = []
+        for item in data:
+            current_text = item.text
+            for fn in self.processing_fn:
+                current_text = fn(current_text)
+            processed_text.append(current_text)
+            
+        embeddings = [
+            self.embedding_fn(text)
+            for text in processed_text
+        ]
+        
+        self.client.insert(self.collection_name, [
+            {
+                "text": item.text,
+                "metadata": item.metadata,
+                "id": item.id,
+                "vector": embedding
+            }
+            for item, embedding in zip(data, embeddings)
+        ])
+        return self
+    
+    def query(self, query: str, output_fields: list[str] = ['text', 'metadata']) -> list[list[dict]]:
+        if not self.embedding_fn:
+            raise ValueError("Embedding function not set. Please call add_embedding_fn first.")
+            
+        preprocessed_query = query
+        for fn in self.processing_fn:
+            preprocessed_query = fn(preprocessed_query)
+        
+        embedding = self.embedding_fn(preprocessed_query)
+        
+        return self.client.search(
+            collection_name=self.collection_name,
+            data=[embedding],
+            limit=10,
+            output_fields=output_fields
+        )
+    
+    def export(self, folder_path: str) -> Self:
+        try:
+            os.makedirs(folder_path, exist_ok=True)
+            
+            # Save using dill
+            pickle_path = os.path.join(folder_path, "ragx.pkl")
+            with open(pickle_path, "wb") as f:
+                dill.dump(self, f)
+            
+            # Copy database
+            db_dest = os.path.join(folder_path, "milvus.db")
+            shutil.copy(self.db_path, db_dest)
+            
+            return self
+            
+        except Exception as e:
+            logger.error(f"Error in export: {e}")
+            raise
+    
+    @classmethod
+    def load(cls, folder_path: str) -> 'Ragx':
+        try:
+            pickle_path = os.path.join(folder_path, "ragx.pkl")
+            
+            with open(pickle_path, "rb") as f:
+                instance = dill.load(f)
+            
+            # Restore client
+            instance.client = MilvusClient(os.path.join(folder_path, "milvus.db"))
+            
+            return instance
+            
+        except Exception as e:
+            logger.error(f"Error in load: {e}")
+            raise
+    
+    def generate_llm_response(self, query: str) -> ChatCompletion:
+        if not self.system_prompt:
+            raise ValueError("System prompt not set. Please call add_system_prompt first.")
+        
+        response = openai.chat.completions.create(
+            model=self.model,
+            messages=[
+                {"role": "system", "content": self.system_prompt},
+                {"role": "user", "content": query}
+            ]
+        )
+        
+        return response