pirag 0.1.7__tar.gz → 0.2.2__tar.gz

{pirag-0.1.7 → pirag-0.2.2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: pirag
-Version: 0.1.7
+Version: 0.2.2
 Summary: CLI Projects of On-Premise RAG. You can use your own LLM and vector DB. Or just add remote LLM servers and vector DB.
 Author-email: semir4in <semir4in@gmail.com>, jyje <jyjeon@outlook.com>
 Project-URL: Homepage, https://github.com/jyje/pilot-onpremise-rag
@@ -9,12 +9,15 @@ Project-URL: Issue, https://github.com/jyje/pilot-onpremise-rag/issues
 Requires-Python: >=3.9
 Description-Content-Type: text/markdown
 License-File: LICENSE
-Requires-Dist: python-dotenv<1.2
+Requires-Dist: dynaconf<3.3
 Requires-Dist: loguru<0.8
 Requires-Dist: pytest<8.4
-Requires-Dist: black<25.2
+Requires-Dist: fastapi<0.116
+Requires-Dist: uvicorn<0.35
 Requires-Dist: ragas<0.3
 Requires-Dist: pymilvus<2.6
+Requires-Dist: langchain-openai<0.4
+Requires-Dist: langchain-ollama<0.4
 Dynamic: license-file
 
 <div align="center">

pirag-0.2.2/app/main.py

@@ -0,0 +1,63 @@
+from loguru import logger
+from argparse import ArgumentParser, ArgumentDefaultsHelpFormatter
+
+import app.rag.config as cfn
+import app.rag.api as api
+import app.rag.cli as cli
+
+# Main parser
+parser = ArgumentParser(
+    formatter_class = ArgumentDefaultsHelpFormatter,
+    description = """
+        Pilot of On-Premise RAG.
+    """,
+    parents = [cfn.top_parser, cfn.common_parser],
+    add_help = False,
+)
+
+# Command definitions
+commands = {
+    # name: help, description, function, extra_parsers
+    "serve"  : ("Start the RAG server",     "Run a FastAPI-based RAG server",              api.serve,  []),
+    "chat"   : ("Chat with the RAG system", "Run an interactive chat with the RAG system", cli.chat,   [cfn.chat_parser]),
+    "train"  : ("Train the RAG system",     "Run a pipeline to train the RAG system",      cli.train,  []),
+    "test"   : ("Test the RAG system",      "Run a pipeline to test the RAG system",       cli.test,   []),
+    "doctor" : ("Diagnose the RAG system",  "Run a pipeline to diagnose the RAG system",   cli.doctor, [cfn.doctor_parser]),
+}
+
+# Add command parsers
+subparsers = parser.add_subparsers(title="commands", dest="command")
+for name, (help, description, _, extra_parsers) in commands.items():
+    subparsers.add_parser(
+        name = name,
+        help = help,
+        description = description,
+        parents = [cfn.common_parser] + extra_parsers,
+        add_help = False,
+    )
+
+def main():
+    args = parser.parse_args()
+    cfn.setup_logger(cfn.LOG_LEVEL, cfn.LOG_SAVE, cfn.LOG_DIR)
+    logger.debug(f"Parsed arguments: {args}")
+
+    if command_info := commands.get(args.command):
+        func, extra_parsers = command_info[2], command_info[3]
+        
+        # Create parser options dict from extra_parsers
+        extra_options = {}
+        if extra_parsers:
+            for parser_obj in extra_parsers:
+                for action in parser_obj._actions:
+                    if action.dest == 'help':
+                        continue
+                    if hasattr(args, action.dest) and getattr(args, action.dest) != action.default:
+                        extra_options[action.dest] = getattr(args, action.dest)
+        
+        # Run the command with the extra parser options
+        func(extra_options)
+    else:
+        parser.print_help()
+
+if __name__ == "__main__":
+    main()

pirag-0.2.2/app/rag/agent/services.py

@@ -0,0 +1,11 @@
+from app.rag.llm.client import client as llm_client
+
+def chat_only_llm():
+    response = llm_client.generate_with_metrics("Hello, how are you?")
+    print(response)
+
+
+def chat_with_rag():
+    pass
+
+

pirag-0.2.2/app/rag/api.py

@@ -0,0 +1,40 @@
+import uvicorn
+from fastapi import FastAPI, APIRouter, Request, Depends, HTTPException, Query
+from fastapi.middleware.cors import CORSMiddleware
+
+from loguru import logger
+import app.rag.config as cfn
+
+from app.rag.routers import system_router
+from app.rag.v1.routers import router as v1_router
+
+# Initialize FastAPI app
+api = FastAPI(
+    title = "RAG API",
+    description = "API for Retrieval-Augmented Generation",
+    version = cfn.__version__,
+)
+
+# Add CORS middleware
+api.add_middleware(
+    CORSMiddleware,
+    allow_origins=["*"],  # Adjust in production
+    allow_credentials=True,
+    allow_methods=["*"],
+    allow_headers=["*"],
+)
+
+api.include_router(router=system_router, prefix="", tags=["System"])
+api.include_router(router=v1_router, prefix="/v1")
+
+def serve(parser_options=None):
+    print("Serving the RAG API...")
+    if parser_options:
+        logger.debug(f"Serve parser options: {parser_options}")
+        
+    uvicorn.run(
+        app = "app.rag.api:api",
+        host = cfn.API_HOST,
+        port = cfn.API_PORT,
+        reload = cfn.API_RELOAD,
+    )

pirag-0.2.2/app/rag/cli.py

@@ -0,0 +1,54 @@
+import app.rag.config as cfn
+from loguru import logger
+
+from app.rag.llm.services import doctor as doctor_llm
+from app.rag.embedding.services import doctor as doctor_embedding
+from app.rag.vector_store.services import doctor as doctor_vector_store
+from app.rag.agent.services import chat_only_llm, chat_with_rag
+
+def chat(options: dict):
+    logger.debug(f"Chat parser options: {options}")
+    no_rag = options.get('no_rag', False)
+
+    # -- Chat
+    if no_rag:
+        logger.info("💬 Chatting with the LLM system directly...")
+        chat_only_llm()
+    else:
+        logger.info("💬 Chatting with the RAG system...")
+        chat_with_rag()
+
+
+def train(options: dict):
+    print("Training the RAG system...")
+    logger.debug(f"Train parser options: {options}")
+
+
+def test(options: dict):
+    print("Testing the RAG system...")
+    logger.debug(f"Test parser options: {options}")
+
+
+def doctor(options: dict):
+    logger.info("💚 Doctoring the RAG system...")
+    
+    logger.debug(f"Doctor parser options: {options}")
+    # Check if resolve option is present
+    resolve = options.get('resolve', False)
+    if resolve:
+        logger.info("🔧 Resolving issues is enabled")
+
+    # -- LLM Server
+    logger.info("🔍 Checking the LLM server (OpenAI-compatible)...")
+    doctor_llm(resolve)
+
+    # -- Embedding Server
+    logger.info("🔍 Checking the embedding server (OpenAI-compatible)...")
+    doctor_embedding(resolve)
+
+    # -- Vector Store
+    logger.info("🔍 Checking the vector store server (Milvus)...")    
+    doctor_vector_store(resolve)
+
+    if resolve:
+        logger.info(f"🔧 Resolving issue completed. To make sure the issues are resolved, please try doctoring again.")

pirag-0.2.2/app/rag/config.py

@@ -0,0 +1,144 @@
+import argparse, sys, pathlib
+from loguru import logger
+from dynaconf import Dynaconf
+from importlib.metadata import version, PackageNotFoundError
+try:
+    __version__ = version("pirag")
+except PackageNotFoundError:
+    __version__ = "0.0.0"
+
+
+# -- Load configuration
+settings = Dynaconf(
+    settings_files = ["settings.yaml"],
+    envvar_prefix = False,
+    load_dotenv = False,
+)
+
+
+# -- Loging
+LOG_LEVEL: str = settings.get("LOG.LEVEL", "INFO").upper()
+if LOG_LEVEL not in ["INFO", "DEBUG", "WARNING", "ERROR", "CRITICAL"]:
+    raise ValueError(f"Invalid log level: {LOG_LEVEL}. Must be one of: INFO, DEBUG, WARNING, ERROR, CRITICAL")
+
+LOG_SAVE: bool = settings.get("LOG.SAVE", False)
+LOG_DIR: str = settings.get("LOG.DIR", ".pirag/logs")
+
+LOG_TIME_FORMAT = "{time:YYYY-MM-DD HH:mm:ss.SSS!UTC}Z"
+LOG_FILE_FORMAT = f"{LOG_TIME_FORMAT} | {{level: <8}} | {{name}}:{{function}}:{{line}} - {{message}}"
+LOG_CONSOLE_FORMAT_FULL = f"<green>{LOG_TIME_FORMAT}</green> | <level>{{level: <8}}</level> | <cyan>{{name}}</cyan>:<cyan>{{function}}</cyan>:<cyan>{{line}}</cyan> - <level>{{message}}</level>\n"
+LOG_CONSOLE_FORMAT_SIMPLE = f"<green>{LOG_TIME_FORMAT}</green> | <level>{{level: <8}}</level> | <level>{{message}}</level>\n"
+
+
+# -- Serving API
+API_HOST: str = settings.get("API.HOST", "0.0.0.0")
+API_PORT: int = settings.get("API.PORT", 8000)
+API_RELOAD: bool = settings.get("API.RELOAD", True)
+
+
+# -- LLM Server
+LLM_BASE_URL: str = settings.get("LLM.BASE_URL", "http://localhost:11434")
+LLM_API_KEY: str = settings.get("LLM.API_KEY", "llm_api_key")
+LLM_MODEL: str = settings.get("LLM.MODEL", "gemma3:4b")
+LLM_SERVER_TYPE: str = settings.get("LLM.SERVER_TYPE", "openai")
+
+
+# -- Embedding Server
+EMBEDDING_BASE_URL: str = settings.get("EMBEDDING.BASE_URL", "http://localhost:11434")
+EMBEDDING_API_KEY: str = settings.get("EMBEDDING.API_KEY", "embedding_api_key")
+EMBEDDING_MODEL: str = settings.get("EMBEDDING.MODEL", "nomic-embed-text:latest")
+EMBEDDING_SERVER_TYPE: str = settings.get("EMBEDDING.SERVER_TYPE", "openai")
+EMBEDDING_DIMENSION: int = settings.get("EMBEDDING.DIMENSION", 768)
+
+
+# -- Data Warehouse
+MINIO_BASE_URL: str = settings.get("MINIO.BASE_URL", "http://localhost:9000")
+MINIO_ACCESS_KEY: str = settings.get("MINIO.ACCESS_KEY", "minioadmin")
+MINIO_SECRET_KEY: str = settings.get("MINIO.SECRET_KEY", "minioadmin")
+MINIO_BUCKET: str = settings.get("MINIO.BUCKET", "pirag")
+MINIO_REGION: str = settings.get("MINIO.REGION", "us-east-1")
+
+
+# -- Vector Store
+MILVUS_BASE_URL: str = settings.get("MILVUS.BASE_URL", "http://localhost:19530")
+MILVUS_USER: str = settings.get("MILVUS.USER", "milvus")
+MILVUS_PASSWORD: str = settings.get("MILVUS.PASSWORD", "milvus")
+MILVUS_DATABASE: str = settings.get("MILVUS.DATABASE", "milvus_database")
+MILVUS_COLLECTION: str = settings.get("MILVUS.COLLECTION", "milvus_collection")
+MILVUS_METRIC_TYPE: str = settings.get("MILVUS.METRIC_TYPE", "IP")
+
+
+# -- Monitoring
+LANGFUSE_BASE_URL: str = settings.get("LANGFUSE.BASE_URL", "http://localhost:8000")
+LANGFUSE_API_KEY: str = settings.get("LANGFUSE.API_KEY", "langfuse_api_key")
+LANGFUSE_PROJECT_ID: str = settings.get("LANGFUSE.PROJECT_ID", "langfuse_project_id")
+
+
+def setup_logger(log_level: str, log_save: bool, log_dir: str):
+    """Configure logger with specified level and outputs"""
+
+    logger.remove()
+
+    # Console handler
+    logger.add(
+        sink = sys.stderr,
+        level = log_level,
+        format = lambda record: LOG_CONSOLE_FORMAT_SIMPLE if record["level"].name == "INFO" else LOG_CONSOLE_FORMAT_FULL,
+        colorize = True
+    )
+
+    if log_save:
+        log_dir = pathlib.Path(log_dir)
+        log_dir.mkdir(exist_ok=True, parents=True)
+
+        # File handler
+        logger.add(
+            sink = log_dir / "{time:YYYYMMDD-HHmmss!UTC}Z.log",
+            level = log_level,
+            rotation = "100 MB",
+            retention = 0,
+            format = LOG_FILE_FORMAT,
+            serialize = False,
+            enqueue = True,
+            backtrace = True,
+            diagnose = True,
+            catch = True
+        )
+
+
+# Top-level parser 
+top_parser = argparse.ArgumentParser(add_help=False)
+top_parser.add_argument(
+    "-v", "--version",
+    help = "Show the `pirag` application's version and exit",
+    action = "version",
+    version = f"{__version__}",
+)
+
+
+# Common parser
+common_parser = argparse.ArgumentParser(add_help=False)
+common_parser.add_argument(
+    "-h", "--help",
+    help = "Show help message and exit",
+    default = argparse.SUPPRESS,
+    action = "help",
+)
+
+
+# Chat parser
+chat_parser = argparse.ArgumentParser(add_help=False)
+chat_parser.add_argument(
+    "-n", "--no-rag",
+    help = "Do not use RAG to answer the question. Just use the LLM to answer the question.",
+    action = "store_true",
+)
+
+
+# Doctor parser
+doctor_parser = argparse.ArgumentParser(add_help=False)
+doctor_parser.add_argument(
+    "-r", "--resolve",
+    help = "Resolve the issue",
+    action = "store_true",
+)

pirag-0.2.2/app/rag/embedding/client.py

@@ -0,0 +1,70 @@
+import requests
+from langchain_openai.embeddings import OpenAIEmbeddings
+
+import app.rag.config as cfn
+from app.rag.utilities import connection_check
+
+
+class EmbeddingClient:
+    def __init__(self, base_url: str, api_key: str, model: str):
+        self.base_url = base_url
+        self.api_key = api_key
+        self.model = model
+        self._is_connected = True
+        self._client = None
+        
+        if self.check_connection():
+            try:
+                self._client = OpenAIEmbeddings(
+                    base_url = base_url,
+                    api_key = api_key,
+                    model = model
+                )
+            except Exception as e:
+                self._is_connected = False
+    
+    def check_connection(self) -> bool:
+        """Check if the embedding server is accessible"""
+        try:
+            requests.head(url=self.base_url, timeout=5)
+        except requests.exceptions.ConnectionError:
+            self._is_connected = False
+            return False
+        self._is_connected = True
+        return True
+    
+    @connection_check
+    def generate(self, prompt: str) -> str:
+        """Generate text from prompt"""
+        if not self._is_connected or self._client is None:
+            return ""
+        return self._client.embed_query(prompt)
+    
+    @connection_check
+    def list_models(self) -> list:
+        """List available models"""
+        if not self._is_connected:
+            return []
+        try:
+            response = requests.get(
+                f"{self.base_url}/models",
+                headers={"Authorization": f"Bearer {self.api_key}"}
+            )
+            if response.status_code == 200:
+                return [model['id'] for model in response.json()['data']]
+            return []
+        except Exception:
+            return []
+    
+    @connection_check
+    def has_model(self, model: str) -> bool:
+        """Check if model exists"""
+        if not self._is_connected:
+            return False
+        return model in self.list_models()
+
+client = EmbeddingClient(
+    base_url = cfn.EMBEDDING_BASE_URL,
+    api_key = cfn.EMBEDDING_API_KEY,
+    model = cfn.EMBEDDING_MODEL,
+)

pirag-0.2.2/app/rag/embedding/services.py

@@ -0,0 +1,26 @@
+from loguru import logger
+
+import app.rag.config as cfn
+from .client import client
+
+def doctor(resolve: bool):
+    # Check connection
+    is_connected = client.check_connection()
+    if not is_connected:
+        logger.error(f"- ❌ FAILED: Embedding connection ({cfn.EMBEDDING_BASE_URL})")
+    else:
+        logger.info(f"- ✅ PASSED: Embedding connection ({cfn.EMBEDDING_BASE_URL})")
+    
+    # Check model availability
+    try:
+        if not is_connected:
+            logger.warning(f"- ⏭️  SKIPPED: Embedding model (Server is not accessible)")
+        else:
+            # List models
+            models = client.list_models()
+            if cfn.EMBEDDING_MODEL not in models:
+                logger.error(f"- ❌ FAILED: Embedding model not found ({cfn.EMBEDDING_MODEL})")
+            else:
+                logger.info(f"- ✅ PASSED: Embedding model found (Model `{cfn.EMBEDDING_MODEL}` exists)")
+    except Exception as e:
+        logger.error(f"- ❌ FAILED: Embedding model check ({str(e)})")

pirag-0.2.2/app/rag/llm/client.py

@@ -0,0 +1,128 @@
+import requests
+import time
+from langchain_openai.llms import OpenAI
+from typing import Dict, Tuple, Any, List, Optional
+
+import app.rag.config as cfn
+from app.rag.utilities import connection_check
+from .utilities import MetricCallbackHandler
+
+class LLMClient:
+    def __init__(self, base_url: str, api_key: str, model: str):
+        self.base_url = base_url
+        self.api_key = api_key
+        self.model = model
+        self._is_connected = True
+        self._client = None
+        
+        if self.check_connection():
+            try:
+                self._client = OpenAI(
+                    base_url = base_url,
+                    api_key = api_key,
+                    model = model
+                )
+            except Exception as e:
+                self._is_connected = False
+    
+    def check_connection(self) -> bool:
+        """Check if the LLM server is accessible"""
+        try:
+            requests.head(url=self.base_url, timeout=5)
+        except requests.exceptions.ConnectionError:
+            self._is_connected = False
+            return False
+        self._is_connected = True
+        return True
+    
+    @connection_check
+    def generate(self, prompt: str) -> tuple:
+        """Generate text from prompt and return usage information
+        
+        Returns:
+            tuple: (generated_text, usage_info)
+        """
+        if not self._is_connected or self._client is None:
+            return "", {}
+        
+        response = self._client.generate([prompt])
+        return response.generations[0][0].text, response.llm_output
+    
+    @connection_check
+    def generate_with_metrics(self, prompt: str) -> Tuple[str, Dict[str, Any]]:
+        """Generate text with timing and usage metrics
+        
+        Returns:
+            tuple: (generated_text, metrics_info)
+        """
+        if not self._is_connected or self._client is None:
+            return "", {"error": "LLM client not connected"}
+        
+        handler = MetricCallbackHandler()
+        
+        # Create streaming client with callback
+        streaming_client = OpenAI(
+            base_url=self.base_url,
+            api_key=self.api_key,
+            model=self.model,
+            streaming=True,
+            callbacks=[handler]
+        )
+        
+        # Make a single request
+        response = streaming_client.generate([prompt], callbacks=[handler])
+        
+        # Get base metrics from response
+        metrics = {}
+        
+        # Extract token usage from response
+        llm_output = response.llm_output if hasattr(response, 'llm_output') else {}
+        
+        # Check if token_usage exists in the response
+        token_usage = llm_output.get('token_usage', {})
+        if token_usage:
+            # If token_usage is available, copy it to our metrics
+            metrics.update(token_usage)
+        
+        # Add model name if available
+        if 'model_name' in llm_output:
+            metrics['model'] = llm_output['model_name']
+        else:
+            metrics['model'] = self.model
+        
+        # Calculate and add timing metrics
+        metrics['ttft'] = handler.ttft or 0.0
+        metrics['total_time'] = (handler.end_time or time.time()) - handler.start_time
+        metrics['tokens_per_second'] = handler.calculate_tokens_per_second()
+        metrics['completion_tokens'] = handler.token_count
+        
+        return handler.result, metrics
+    
+    @connection_check
+    def list_models(self) -> list:
+        """List available models"""
+        if not self._is_connected:
+            return []
+        try:
+            response = requests.get(
+                f"{self.base_url}/models",
+                headers={"Authorization": f"Bearer {self.api_key}"}
+            )
+            if response.status_code == 200:
+                return [model['id'] for model in response.json()['data']]
+            return []
+        except Exception:
+            return []
+    
+    @connection_check
+    def has_model(self, model: str) -> bool:
+        """Check if model exists"""
+        if not self._is_connected:
+            return False
+        return model in self.list_models()
+
+client = LLMClient(
+    base_url = cfn.LLM_BASE_URL,
+    api_key = cfn.LLM_API_KEY,
+    model = cfn.LLM_MODEL,
+)

pirag-0.2.2/app/rag/llm/services.py

@@ -0,0 +1,26 @@
+from loguru import logger
+
+import app.rag.config as cfn
+from .client import client
+
+def doctor(resolve: bool):
+    # Check connection
+    is_connected = client.check_connection()
+    if not is_connected:
+        logger.error(f"- ❌ FAILED: LLM connection ({cfn.LLM_BASE_URL})")
+    else:
+        logger.info(f"- ✅ PASSED: LLM connection ({cfn.LLM_BASE_URL})")
+    
+    # Check model availability
+    try:
+        if not is_connected:
+            logger.warning(f"- ⏭️  SKIPPED: LLM model (Server is not accessible)")
+        else:
+            # List models
+            models = client.list_models()
+            if cfn.LLM_MODEL not in models:
+                logger.error(f"- ❌ FAILED: LLM model not found ({cfn.LLM_MODEL})")
+            else:
+                logger.info(f"- ✅ PASSED: LLM model found (Model `{cfn.LLM_MODEL}` exists)")
+    except Exception as e:
+        logger.error(f"- ❌ FAILED: LLM model check ({str(e)})")

pirag-0.2.2/app/rag/llm/utilities.py

@@ -0,0 +1,40 @@
+import time
+from langchain.callbacks.base import BaseCallbackHandler
+
+class MetricCallbackHandler(BaseCallbackHandler):
+    def __init__(self):
+        self.start_time = time.time()
+        self.ttft = None
+        self.first_token_time = None
+        self.result = ""
+        self.end_time = None
+        self.token_count = 0
+        self.token_timestamps = []
+    
+    def on_llm_new_token(self, token: str, **kwargs):
+        current_time = time.time()
+        self.token_count += 1
+        self.token_timestamps.append(current_time)
+        
+        if self.ttft is None:
+            self.ttft = current_time - self.start_time
+            self.first_token_time = current_time
+            
+        self.result += token
+    
+    def on_llm_end(self, *args, **kwargs):
+        self.end_time = time.time()
+    
+    def calculate_tokens_per_second(self):
+        """Calculate tokens per second after the first token"""
+        if self.token_count <= 1 or self.first_token_time is None or self.end_time is None:
+            return 0.0
+        
+        # Calculate time from first token to completion (exclude TTFT)
+        generation_time = self.end_time - self.first_token_time
+        if generation_time <= 0:
+            return 0.0
+        
+        # Exclude the first token from the count since we're measuring from after it arrived
+        tokens_after_first = self.token_count - 1
+        return tokens_after_first / generation_time

pirag-0.2.2/app/rag/models.py

@@ -0,0 +1,19 @@
+from pydantic import BaseModel
+
+class SystemStatusResponse(BaseModel):
+    """
+    Response model for the system status endpoint.
+    """
+    status: int
+    message: str
+    
+    model_config = {
+        "json_schema_extra": {
+            "examples": [
+                {
+                    "status": 200,
+                    "message": "System is running normally"
+                }
+            ]
+        }
+    }