pirag 0.1.7__py3-none-any.whl → 0.2.2__py3-none-any.whl

app/main.py

@@ -1,81 +1,63 @@
-import argparse, os
-from dotenv import load_dotenv
 from loguru import logger
+from argparse import ArgumentParser, ArgumentDefaultsHelpFormatter
 
-load_dotenv(dotenv_path=os.environ.get('ENV_FILE', '.env'), override=True)
-
-from app.rag.config import top_parser, common_parser, setup_logger
-from app.rag.doctor import help as doctor_help, parser as doctor_parser, route as doctor_route
-from app.rag.train import help as train_help, parser as train_parser, route as train_route
-from app.rag.ask import help as ask_help, parser as ask_parser, route as ask_route
-from app.rag.test import help as test_help, parser as test_parser, route as test_route
+import app.rag.config as cfn
+import app.rag.api as api
+import app.rag.cli as cli
 
 # Main parser
-parser = argparse.ArgumentParser(
-    formatter_class = argparse.ArgumentDefaultsHelpFormatter,
-    description = 'Pilot of On-Premise RAG',
-    parents = [top_parser],
-    add_help = False,
-)
-
-# Commands
-subparsers = parser.add_subparsers(
-    title = 'commands',
-    dest = 'command',
-)
-
-subparsers.add_parser(
-    'doctor',
-    help = doctor_help,
-    description = doctor_parser.description,
-    parents = [top_parser, common_parser, doctor_parser],
-    add_help = False,
-)
-
-subparsers.add_parser(
-    'train',
-    help = train_help,
-    description = train_parser.description,
-    parents = [top_parser, common_parser, train_parser],
-    add_help = False,
-)
-
-subparsers.add_parser(
-    'test',
-    help = test_help,
-    description = test_parser.description,
-    parents = [top_parser, common_parser, test_parser],
+parser = ArgumentParser(
+    formatter_class = ArgumentDefaultsHelpFormatter,
+    description = """
+        Pilot of On-Premise RAG.
+    """,
+    parents = [cfn.top_parser, cfn.common_parser],
     add_help = False,
 )
 
-subparsers.add_parser(
-    'ask',
-    help = ask_help,
-    description = ask_parser.description,
-    parents = [top_parser, common_parser, ask_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()
-    setup_logger(
-        log_level = args.log_level,
-        log_dir = args.log_dir,
-    )
-    command_message = f"with command: {args.command}" if args.command else ""
-    logger.info(f"RAG Started {command_message}")
+    cfn.setup_logger(cfn.LOG_LEVEL, cfn.LOG_SAVE, cfn.LOG_DIR)
     logger.debug(f"Parsed arguments: {args}")
 
-    if args.command == 'doctor':
-        doctor_route(args)
-    elif args.command == 'ask':
-        ask_route(args)
-    elif args.command == 'train':
-        train_route(args)
-    elif args.command == 'test':
-        test_route(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__':
+if __name__ == "__main__":
     main()

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
+
+

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,
+    )

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.")

app/rag/config.py

@@ -1,35 +1,84 @@
-import argparse, os, sys
-from pathlib import Path
+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")
 
-# Logger format constants
 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"
 
-def setup_logger(log_level: str, log_dir: str):
+
+# -- 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"""
-    
-    log_dir = Path(log_dir)
-    log_dir.mkdir(exist_ok=True, parents=True)
-    
+
     logger.remove()
-    
-    # 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
-    )
-    
+
     # Console handler
     logger.add(
         sink = sys.stderr,
@@ -38,114 +87,58 @@ def setup_logger(log_level: str, log_dir: str):
         colorize = True
     )
 
-
-class EnvDefault(argparse.Action):
-    """Custom argparse action that uses environment variables as defaults.
-
-    This action extends the standard argparse.Action to support reading default values
-    from environment variables. If the specified environment variable exists, its value
-    will be used as the default value for the argument.
-
-    For boolean flags (store_true/store_false), the environment variable is interpreted
-    as a boolean value where 'true', '1', 'yes', or 'on' (case-insensitive) are
-    considered True.
-
-    Args:
-        envvar (str): Name of the environment variable to use as default
-        required (bool, optional): Whether the argument is required. Defaults to True.
-            Note: If a default value is found in environment variables, required is set to False.
-        default (Any, optional): Default value if environment variable is not set. Defaults to None.
-        **kwargs: Additional arguments passed to argparse.Action
-
-    Example:
-        ```python
-        parser.add_argument(
-            '--log-level',
-            envvar='LOG_LEVEL',
-            help='Logging level',
-            default='INFO',
-            action=EnvDefault
+    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
         )
-        ```
-
-    Note:
-        The help text is automatically updated to include the environment variable name.
-    """
-    def __init__(self, envvar, required=True, default=None, **kwargs):
-        if envvar and envvar in os.environ:
-            env_value = os.environ[envvar]
-            # Convert string environment variable to boolean
-            if kwargs.get('nargs') is None and kwargs.get('const') is not None:  # store_true/store_false case
-                default = env_value.lower() in ('true', '1', 'yes', 'on')
-            else:
-                default = env_value
-            logger.debug(f"Using {envvar}={default} from environment")
-        
-        if envvar:
-            kwargs["help"] += f" (envvar: {envvar})"
-        
-        if required and default:
-            required = False
-            
-        super(EnvDefault, self).__init__(default=default, required=required, **kwargs)
-        self.envvar = envvar
-
-    def __call__(self, parser, namespace, values, option_string=None):
-        setattr(namespace, self.dest, values if values is not None else self.default)
-
-
-# Top-level parser with common options
-top_parser = argparse.ArgumentParser(add_help=False)
 
-top_parser.add_argument(
-    '-h', '--help',
-    help = 'Show help message and exit',
-    default = argparse.SUPPRESS,
-    action = 'help',
-)
 
+# Top-level parser 
+top_parser = argparse.ArgumentParser(add_help=False)
 top_parser.add_argument(
-    '--env-file',
-    envvar = 'ENV_FILE',
-    help = 'Path to environment file',
-    default = '.env',
-    type = str,
-    action = EnvDefault,
+    "-v", "--version",
+    help = "Show the `pirag` application's version and exit",
+    action = "version",
+    version = f"{__version__}",
 )
 
-top_parser.add_argument(
-    '--log-level',
-    envvar = 'LOG_LEVEL',
-    help = 'Logging level',
-    default = 'INFO',
-    type = lambda x: x.upper(),
-    choices = ['DEBUG', 'INFO', 'WARNING', 'ERROR', 'CRITICAL'],
-    required = False,
-    action = EnvDefault,
-)
 
-top_parser.add_argument(
-    '--log-dir',
-    envvar = 'LOG_DIR',
-    help = 'Path to log directory',
-    default = '.pirag/logs',
-    type = str,
-    required = False,
-    action = EnvDefault,
+# 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",
 )
 
-top_parser.add_argument(
-    '--log-save',
-    envvar = 'LOG_SAVE',
-    help = 'Save log to file. If this flag is set, the log will be saved to the file specified in the `--log-path`.',
-    default = False,
-    const = True,
-    nargs = 0,
-    type = bool,
-    required = False,
-    action = EnvDefault,
+
+# 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",
 )
 
-common_parser = argparse.ArgumentParser(
-    add_help = False,
+
+# Doctor parser
+doctor_parser = argparse.ArgumentParser(add_help=False)
+doctor_parser.add_argument(
+    "-r", "--resolve",
+    help = "Resolve the issue",
+    action = "store_true",
 )

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,
+)

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)})")