ragxo 0.1.3__py3-none-any.whl → 0.1.5__py3-none-any.whl

ragxo/__init__.py

@@ -1,3 +1,3 @@
-from ragx import Ragxo, Document
 
+from .client import Ragxo, Document
 __all__ = ["Ragxo", "Document"]

ragxo/client.py

@@ -0,0 +1,253 @@
+import time
+from typing import Self, Callable
+from pymilvus import MilvusClient
+from pydantic import BaseModel
+import boto3
+import dill
+import os
+import shutil
+import logging
+import tempfile
+from botocore.exceptions import ClientError
+import openai
+from openai import ChatCompletion
+
+logger = logging.getLogger(__name__)
+
+class Document(BaseModel):
+    text: str
+    metadata: dict
+    id: int
+
+class Ragxo:
+    def __init__(self, dimension: int) -> None:
+        self.dimension = dimension
+        self.collection_name = "ragx"
+        os.makedirs("ragx_artifacts", exist_ok=True)
+
+        self.db_path = f"ragx_artifacts/milvus_{int(time.time())}.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, destination: str, s3_bucket: str = None) -> Self:
+        """
+        Export the Ragx instance to either local filesystem or S3.
+        
+        Args:
+            destination: str - Local path or S3 key prefix
+            s3_bucket: str, optional - S3 bucket name. If provided, export to S3
+        """
+        try:
+            # If s3_bucket is provided, export to S3
+            if s3_bucket:
+                return self._export_to_s3(destination, s3_bucket)
+            
+            # Otherwise, export to local filesystem
+            os.makedirs(destination, exist_ok=True)
+            
+            # Save using dill
+            pickle_path = os.path.join(destination, "ragx.pkl")
+            with open(pickle_path, "wb") as f:
+                dill.dump(self, f)
+            
+            # Copy database
+            db_dest = os.path.join(destination, "milvus.db")
+            shutil.copy(self.db_path, db_dest)
+            
+            return self
+            
+        except Exception as e:
+            logger.error(f"Error in export: {e}")
+            raise
+
+    def _export_to_s3(self, prefix: str, bucket: str) -> Self:
+        """
+        Internal method to handle S3 export.
+        """
+        try:
+            s3_client = boto3.client('s3')
+            
+            # Create a temporary directory for the files
+            with tempfile.TemporaryDirectory() as temp_dir:
+                # Save pickle file
+                pickle_path = os.path.join(temp_dir, "ragx.pkl")
+                with open(pickle_path, "wb") as f:
+                    dill.dump(self, f)
+                
+                # Copy database
+                db_path = os.path.join(temp_dir, "milvus.db")
+                shutil.copy(self.db_path, db_path)
+                
+                # Upload to S3
+                s3_client.upload_file(
+                    pickle_path,
+                    bucket,
+                    f"{prefix}/ragx.pkl"
+                )
+                s3_client.upload_file(
+                    db_path,
+                    bucket,
+                    f"{prefix}/milvus.db"
+                )
+            
+            return self
+            
+        except ClientError as e:
+            logger.error(f"Error uploading to S3: {e}")
+            raise
+        except Exception as e:
+            logger.error(f"Error in S3 export: {e}")
+            raise
+
+    @classmethod
+    def load(cls, source: str, s3_bucket: str = None) -> Self:
+        """
+        Load a Ragx instance from either local filesystem or S3.
+        
+        Args:
+            source: str - Local path or S3 key prefix
+            s3_bucket: str, optional - S3 bucket name. If provided, load from S3
+        """
+        try:
+            # If s3_bucket is provided, load from S3
+            if s3_bucket:
+                return cls._load_from_s3(source, s3_bucket)
+            
+            # Otherwise, load from local filesystem
+            pickle_path = os.path.join(source, "ragx.pkl")
+            
+            with open(pickle_path, "rb") as f:
+                instance = dill.load(f)
+            
+            # Restore client
+            instance.client = MilvusClient(os.path.join(source, "milvus.db"))
+            
+            return instance
+            
+        except Exception as e:
+            logger.error(f"Error in load: {e}")
+            raise
+
+    @classmethod
+    def _load_from_s3(cls, prefix: str, bucket: str) -> 'Ragx':
+        """
+        Internal classmethod to handle S3 loading.
+        """
+        try:
+            s3_client = boto3.client('s3')
+            
+            # Create a temporary directory for the files
+            with tempfile.TemporaryDirectory() as temp_dir:
+                # Download files from S3
+                pickle_path = os.path.join(temp_dir, "ragx.pkl")
+                db_path = os.path.join(temp_dir, "milvus.db")
+                
+                s3_client.download_file(
+                    bucket,
+                    f"{prefix}/ragx.pkl",
+                    pickle_path
+                )
+                s3_client.download_file(
+                    bucket,
+                    f"{prefix}/milvus.db",
+                    db_path
+                )
+                
+                # Load the pickle file
+                with open(pickle_path, "rb") as f:
+                    instance = dill.load(f)
+                
+                # Restore client with the downloaded database
+                instance.client = MilvusClient(db_path)
+                
+                return instance
+                
+        except ClientError as e:
+            logger.error(f"Error downloading from S3: {e}")
+            raise
+        except Exception as e:
+            logger.error(f"Error in S3 load: {e}")
+            raise
+    
+    def generate_llm_response(self, query: str, data: list[dict] = None) -> ChatCompletion:
+        if data is None:
+            data = self.query(query)[0]
+        
+        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: {} data: {}".format(query, data)}
+            ]
+        )
+        
+        return response

{ragxo-0.1.3.dist-info → ragxo-0.1.5.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: ragxo
-Version: 0.1.3
+Version: 0.1.5
 Summary: A RAG (Retrieval-Augmented Generation) toolkit with Milvus integration
 Home-page: https://github.com/yourusername/ragx
 License: MIT
@@ -17,6 +17,7 @@ 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: boto3 (>=1.36.14,<2.0.0)
 Requires-Dist: dill (>=0.3.9,<0.4.0)
 Requires-Dist: milvus (>=2.3.9,<3.0.0)
 Requires-Dist: openai (>=1.61.1,<2.0.0)
@@ -48,65 +49,15 @@ RagXO extends the capabilities of traditional RAG (Retrieval-Augmented Generatio
 pip install ragxo
 ```
 
-## Quick Start 🚀
-
-```python
-from ragxo import Ragxo, Document
-from openai import OpenAI
-client = OpenAI()
-
-def get_openai_embeddings(text: str) -> list[float]:
-    response = client.embeddings.create(
-        input=text,
-        model="text-embedding-ada-002"
-    )
-    return response.data[0].embedding
-
-def preprocess_text(text: str) -> str:
-    return text.lower()
-
-# Initialize and configure RagXO
-ragxo = Ragxo(dimension=384)
-ragxo.add_preprocess(preprocess_text)
-ragxo.add_embedding_fn(get_openai_embeddings)
-
-# Add system prompt and model
-ragxo.add_system_prompt("You are a helpful assistant.")
-ragxo.add_model("gpt-4o-mini")
-
-# Create and index documents
-documents = [
-    Document(
-        text="Sample document for indexing",
-        metadata={"source": "example"},
-        id=1
-    )
-]
-ragxo.index(documents)
-
-# Export the pipeline
-ragxo.export("my_rag_v1")
-
-# Load and use elsewhere
-loaded_ragxo = Ragxo.load("my_rag_v1")
-
-# Query and generate response
-similar_docs = loaded_ragxo.query("sample query")
-llm_response = loaded_ragxo.generate_llm_response("What can you tell me about the sample?")
-```
-
 ## Usage Guide 📚
 
-### Creating Documents
+### Import
 
 ```python
-from ragxo import Document
+from ragxo import Ragxo, Document
+
+ragxo_client = Ragxo(dimension=768)
 
-doc = Document(
-    text="Your document content here",
-    metadata={"source": "wiki", "category": "science"},
-    id=1
-)
 ```
 
 ### Adding Preprocessing Steps
@@ -120,8 +71,8 @@ def remove_special_chars(text: str) -> str:
 def lowercase(text: str) -> str:
     return text.lower()
 
-ragxo.add_preprocess(remove_special_chars)
-ragxo.add_preprocess(lowercase)
+ragxo_client.add_preprocess(remove_special_chars)
+ragxo_client.add_preprocess(lowercase)
 ```
 
 ### Custom Embedding Functions
@@ -150,27 +101,43 @@ def get_openai_embeddings(text: str) -> list[float]:
 ragxo.add_embedding_fn(get_openai_embeddings)
 ```
 
+
+### Creating Documents
+
+```python
+from ragxo import Document
+
+doc = Document(
+    text="Your document content here",
+    metadata={"source": "wiki", "category": "science"},
+    id=1
+)
+
+ragxo_client.index([doc])
+
+```
+
 ### LLM Configuration
 
 ```python
 # Set system prompt
-ragxo.add_system_prompt("""
+ragxo_client.add_system_prompt("""
 You are a helpful assistant. Use the provided context to answer questions accurately.
 If you're unsure about something, please say so.
 """)
 
 # Set LLM model
-ragxo.add_model("gpt-4")
+ragxo_client.add_model("gpt-4")
 ```
 
 ### Export and Load
 
 ```python
 # Export your RAG pipeline
-ragxo.export("rag_pipeline_v1")
+ragxo_client.export("rag_pipeline_v1")
 
 # Load it elsewhere
-loaded_ragxo = Ragxo.load("rag_pipeline_v1")
+loaded_ragxo_client = Ragxo.load("rag_pipeline_v1")
 ```
 
 ## Best Practices 💡
@@ -180,27 +147,15 @@ loaded_ragxo = Ragxo.load("rag_pipeline_v1")
 ragxo.export("my_rag_v1.0.0")
 ```
 
-2. **Validate After Loading**: Always test your loaded pipeline:
-```python
-loaded_ragxo = Ragxo.load("my_rag")
-try:
-    # Test similarity search
-    similar_docs = loaded_ragxo.query("test query")
-    # Test LLM generation
-    llm_response = loaded_ragxo.generate_llm_response("test question")
-    print("Pipeline loaded successfully!")
-except Exception as e:
-    print(f"Error loading pipeline: {e}")
+2. **S3**: Use S3 to store your exports
+
+```shell
+export AWS_ACCESS_KEY_ID=your_access_key
+export AWS_SECRET_ACCESS_KEY=your_secret_key
 ```
 
-3. **Document Your Pipeline Configuration**: Keep track of your setup:
 ```python
-pipeline_config = {
-    "preprocessing_steps": ["remove_special_chars", "lowercase"],
-    "embedding_model": "all-MiniLM-L6-v2",
-    "llm_model": "gpt-4",
-    "dimension": 384
-}
+ragxo_client.export("my_rag_v1.0.0", s3_bucket="my_bucket")
 ```
 
 ## License 📝

ragxo-0.1.5.dist-info/RECORD

@@ -0,0 +1,5 @@
+ragxo/__init__.py,sha256=0VVe-z4XkkGQLQIG0hF0Hyf87_RgX0E4T9TRwwTkbmE,68
+ragxo/client.py,sha256=smS3vt7k0k1p1mDOT8Taa4vpHl0fiuvbo3RdRY_D01k,8300
+ragxo-0.1.5.dist-info/METADATA,sha256=0oS9x3tRULcvKZNQQOyIlKpQ_baDzt1pHB0wTubsb7I,4600
+ragxo-0.1.5.dist-info/WHEEL,sha256=Nq82e9rUAnEjt98J6MlVmMCZb-t9cYE2Ir1kpBmnWfs,88
+ragxo-0.1.5.dist-info/RECORD,,

ragxo/ragx.py

@@ -1,145 +0,0 @@
-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 Ragxo:
-    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, data: list[dict] = None) -> ChatCompletion:
-        
-        if data is None:
-            data = self.query(query)[0]
-        
-        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: {} data: {}".format(query, data)}
-            ]
-        )
-        
-        return response

ragxo-0.1.3.dist-info/RECORD

@@ -1,5 +0,0 @@
-ragxo/__init__.py,sha256=jI_6iulTUQk9JUDft-jM6NHESpZSmJVPIaVOmd4-jWw,65
-ragxo/ragx.py,sha256=_HQCTth_iR2rxV9amMyA6qlOpdGji5_-rSDB5WWG2u4,4537
-ragxo-0.1.3.dist-info/METADATA,sha256=FZmy-PL_SZMf9NuDWcniQUsleZna_GYsz5GLoJRbHcM,5960
-ragxo-0.1.3.dist-info/WHEEL,sha256=Nq82e9rUAnEjt98J6MlVmMCZb-t9cYE2Ir1kpBmnWfs,88
-ragxo-0.1.3.dist-info/RECORD,,