ragforge-framework 0.1.0__py3-none-any.whl

ragforge_framework/cache/user_manager.py

@@ -0,0 +1,264 @@
+from hashlib import sha256
+from datetime import datetime
+
+from .repository import Repository
+
+
+class UserManager:
+    """
+    Manages user registration, requests tracking, rate limiting, and analytics.
+    """
+    def __init__(self, config):
+        self.config = config
+        self.repository = Repository(config.USER_CACHE_DIR)
+        self.daily_limit = config.DAILY_LIMIT
+
+    # ==========================================================
+    # PRIVATE HELPERS
+    # ==========================================================
+
+    def _today(self):
+        return datetime.now().strftime("%Y-%m-%d")
+
+    def _now(self):
+        return datetime.now().strftime("%Y-%m-%d %H:%M:%S")
+
+    def _user_id(self, ip: str):
+        return sha256(ip.encode()).hexdigest()
+
+    def _cache_key(self, ip: str):
+        return f"{self._user_id(ip)}:{self._today()}"
+
+    # ==========================================================
+    # USER FACTORY
+    # ==========================================================
+
+    def _create_user(self, ip: str):
+        return {
+            "user_id": self._user_id(ip),
+            "ip": ip,
+            "date": self._today(),
+            "total_requests": 0,
+            "llm_requests": 0,
+            "remaining": self.daily_limit,
+            "first_request": self._now(),
+            "last_request": self._now(),
+            "cache_hits": 0,
+            "cache_misses": 0,
+            "provider_usage": {},
+            "fallbacks": 0
+        }
+
+    # ==========================================================
+    # GET USER
+    # ==========================================================
+
+    def get_user(self, ip: str):
+        key = self._cache_key(ip)
+        user = self.repository.get(key)
+        if user is None:
+            user = self._create_user(ip)
+            self.repository.set(key, user)
+        return user
+
+    # ==========================================================
+    # SAVE USER
+    # ==========================================================
+
+    def save_user(self, user: dict):
+        key = f"{user['user_id']}:{user['date']}"
+        self.repository.set(key, user)
+
+    # ==========================================================
+    # REQUEST COUNTER
+    # ==========================================================
+
+    def register_request(self, ip: str):
+        user = self.get_user(ip)
+        user["total_requests"] += 1
+        user["last_request"] = self._now()
+        self.save_user(user)
+        return user
+
+    # ==========================================================
+    # RATE LIMIT
+    # ==========================================================
+
+    def allow_request(self, ip: str):
+        user = self.get_user(ip)
+        if user["llm_requests"] >= self.daily_limit:
+            return False, user
+        user["llm_requests"] += 1
+        user["remaining"] = self.daily_limit - user["llm_requests"]
+        user["last_request"] = self._now()
+        self.save_user(user)
+        return True, user
+
+    # ==========================================================
+    # CACHE HIT
+    # ==========================================================
+
+    def cache_hit(self, ip: str):
+        user = self.get_user(ip)
+        user["cache_hits"] += 1
+        self.save_user(user)
+
+    # ==========================================================
+    # CACHE MISS
+    # ==========================================================
+
+    def cache_miss(self, ip: str):
+        user = self.get_user(ip)
+        user["cache_misses"] += 1
+        self.save_user(user)
+
+    # ==========================================================
+    # PROVIDER
+    # ==========================================================
+
+    def provider_used(self, ip: str, provider: str):
+        user = self.get_user(ip)
+        provider = provider.lower()
+        if provider not in user["provider_usage"]:
+            user["provider_usage"][provider] = 0
+        user["provider_usage"][provider] += 1
+        self.save_user(user)
+
+    # ==========================================================
+    # FALLBACK
+    # ==========================================================
+
+    def record_fallback(self, ip: str):
+        user = self.get_user(ip)
+        user["fallbacks"] += 1
+        self.save_user(user)
+
+    # ==========================================================
+    # USER EXISTS
+    # ==========================================================
+
+    def user_exists(self, ip: str) -> bool:
+        key = self._cache_key(ip)
+        return self.repository.exists(key)
+
+    # ==========================================================
+    # GET ALL USERS
+    # ==========================================================
+
+    def get_all_users(self):
+        return self.repository.values()
+
+    # ==========================================================
+    # ACTIVE USERS
+    # ==========================================================
+
+    def active_users(self):
+        users = self.get_all_users()
+        return sorted(
+            users,
+            key=lambda user: user["last_request"],
+            reverse=True
+        )
+
+    # ==========================================================
+    # SEARCH USER
+    # ==========================================================
+
+    def search_user(self, user_id: str):
+        users = self.get_all_users()
+        for user in users:
+            if user["user_id"] == user_id:
+                return user
+        return None
+
+    # ==========================================================
+    # RESET USER
+    # ==========================================================
+
+    def reset_user(self, user_id: str):
+        users = self.repository.items()
+        for key, user in users:
+            if user["user_id"] == user_id:
+                self.repository.delete(key)
+                return True
+        return False
+
+    # ==========================================================
+    # RESET ALL USERS
+    # ==========================================================
+
+    def reset_all(self):
+        self.repository.clear()
+
+    # ==========================================================
+    # DELETE OLD USERS
+    # ==========================================================
+
+    def delete_old_users(self):
+        today = self._today()
+        users = self.repository.items()
+        deleted = 0
+        for key, user in users:
+            if user["date"] != today:
+                self.repository.delete(key)
+                deleted += 1
+        return deleted
+
+    # ==========================================================
+    # TOP USERS
+    # ==========================================================
+
+    def get_top_users(self, limit: int = 10):
+        users = self.get_all_users()
+        users.sort(
+            key=lambda user: user["llm_requests"],
+            reverse=True
+        )
+        return users[:limit]
+
+    # ==========================================================
+    # STATISTICS
+    # ==========================================================
+
+    def get_statistics(self):
+        users = self.get_all_users()
+        stats = {
+            "total_users": len(users),
+            "total_requests": 0,
+            "llm_requests": 0,
+            "cache_hits": 0,
+            "cache_misses": 0,
+            "fallbacks": 0,
+            "provider_usage": {}
+        }
+        for user in users:
+            stats["total_requests"] += user["total_requests"]
+            stats["llm_requests"] += user["llm_requests"]
+            stats["cache_hits"] += user["cache_hits"]
+            stats["cache_misses"] += user["cache_misses"]
+            stats["fallbacks"] += user["fallbacks"]
+            for provider, count in user["provider_usage"].items():
+                if provider not in stats["provider_usage"]:
+                    stats["provider_usage"][provider] = 0
+                stats["provider_usage"][provider] += count
+        return stats
+
+    # ==========================================================
+    # EXPORT USERS
+    # ==========================================================
+
+    def export_users(self):
+        return {
+            "users": self.get_all_users(),
+            "statistics": self.get_statistics()
+        }
+
+    # ==========================================================
+    # DEBUG
+    # ==========================================================
+
+    def __len__(self):
+        return self.repository.count()
+
+    def __repr__(self):
+        return f"<UserManager users={len(self)}>"

ragforge_framework/cli.py

@@ -0,0 +1,91 @@
+import argparse
+import sys
+import os
+
+def main():
+    """
+    RAGForge CLI Entry Point.
+    Processes command-line instructions for template initialization, vector ingestion, and API execution.
+    """
+    parser = argparse.ArgumentParser(description="RAGForge Framework Command Line Utility")
+    subparsers = parser.add_subparsers(dest="command", help="Available CLI commands")
+
+    # Command: init
+    subparsers.add_parser("init", help="Initialize a new RAGForge project template")
+
+    # Command: ingest
+    subparsers.add_parser("ingest", help="Run the document chunking and vector indexing pipeline")
+
+    # Command: run
+    run_parser = subparsers.add_parser("run", help="Start the FastAPI uvicorn backend")
+    run_parser.add_argument("--host", default="127.0.0.1", help="Host address to bind uvicorn server")
+    run_parser.add_argument("--port", type=int, default=8000, help="Port to run the backend API")
+
+    args = parser.parse_args()
+
+    if args.command == "init":
+        print("\n[RAGForge] Initializing starter template...")
+        os.makedirs("data", exist_ok=True)
+
+        data_file_path = os.path.join("data", "data.txt")
+        if not os.path.exists(data_file_path):
+            with open(data_file_path, "w", encoding="utf-8") as f:
+                f.write("Yashodeep Hundiwale is an experienced AI Engineer and creator of RAGForge.\n")
+            print(f"  [+] Created default text source: {data_file_path}")
+
+        app_file_path = "app.py"
+        if not os.path.exists(app_file_path):
+            with open(app_file_path, "w", encoding="utf-8") as f:
+                f.write('''from ragforge_framework import RAGForge
+
+# Instantiate the RAGForge framework config and services
+forge = RAGForge(
+    data_path="data/data.txt",
+    faiss_path="faiss_index",
+    primary_model="ollama/llama3.2:latest"
+)
+
+# Expose fastapi app for production WSGI/ASGI servers (e.g. gunicorn/uvicorn)
+app = forge.app
+
+if __name__ == "__main__":
+    # Start the server locally when running app.py directly
+    forge.run(host="127.0.0.1", port=8000)
+''')
+            print(f"  [+] Created default project entry point: {app_file_path}")
+
+        print("\n[RAGForge] Initialization completed successfully!")
+        print("  * Ingest your document database:  ragforge ingest")
+        print("  * Launch your backend API:        ragforge run\n")
+
+    elif args.command == "ingest":
+        print("\n[RAGForge] Loading Ingestion Pipeline...")
+        forge = get_forge_instance()
+        forge.ingest()
+
+    elif args.command == "run":
+        print(f"\n[RAGForge] Starting server on {args.host}:{args.port}...")
+        forge = get_forge_instance()
+        forge.run(host=args.host, port=args.port)
+
+    else:
+        parser.print_help()
+
+
+def get_forge_instance():
+    """
+    Looks for a local 'app.py' file in the current working directory to load
+    custom configs; falls back to default settings if app.py is missing or invalid.
+    """
+    if os.path.exists("app.py"):
+        sys.path.insert(0, os.getcwd())
+        try:
+            import app as local_app
+            if hasattr(local_app, "forge"):
+                return local_app.forge
+        except Exception as e:
+            print(f"  [!] Warning: Failed to load 'forge' instance from local 'app.py' due to: {e}")
+            print("  Falling back to default framework configuration...")
+
+    from ragforge_framework.core import RAGForge
+    return RAGForge()

ragforge_framework/config.py

@@ -0,0 +1,50 @@
+import os
+from typing import List, Optional
+
+class RAGForgeConfig:
+    """
+    RAGForge Configuration class.
+    Stores and validates all framework parameters.
+    """
+    def __init__(
+        self,
+        primary_model: str = "ollama/llama3.2:latest",
+        fallback_models: Optional[List[str]] = None,
+        temperature: float = 0.2,
+        timeout: int = 30,
+        max_retries: int = 3,
+        admin_key: Optional[str] = None,
+        daily_limit: int = 10,
+        response_cache_dir: str = "./response_cache",
+        user_cache_dir: str = "./user_cache",
+        data_path: str = "data/data.txt",
+        faiss_path: str = "faiss_index",
+        embedding_model: str = "sentence-transformers/all-MiniLM-L6-v2",
+        chunk_size: int = 800,
+        chunk_overlap: int = 100,
+        search_type: str = "mmr",
+        top_k: int = 3,
+        fetch_k: int = 5,
+        lambda_mult: float = 0.5,
+    ):
+        self.PRIMARY_MODEL = primary_model
+        self.FALLBACK_MODELS = fallback_models or [
+            "groq/llama-3.1-8b-instant",
+            "mistralai/mistral-small-latest",
+        ]
+        self.TEMPERATURE = temperature
+        self.TIMEOUT = timeout
+        self.MAX_RETRIES = max_retries
+        self.ADMIN_KEY = admin_key or os.getenv("ADMIN_KEY") or "yashodeep"
+        self.DAILY_LIMIT = daily_limit
+        self.RESPONSE_CACHE_DIR = response_cache_dir
+        self.USER_CACHE_DIR = user_cache_dir
+        self.DATA_PATH = data_path
+        self.FAISS_PATH = faiss_path
+        self.EMBEDDING_MODEL = embedding_model
+        self.CHUNK_SIZE = chunk_size
+        self.CHUNK_OVERLAP = chunk_overlap
+        self.SEARCH_TYPE = search_type
+        self.TOP_K = top_k
+        self.FETCH_K = fetch_k
+        self.LAMBDA_MULT = lambda_mult

ragforge_framework/core.py

@@ -0,0 +1,42 @@
+import os
+import uvicorn
+from .config import RAGForgeConfig
+from .app import create_app
+from .rag.injest import ingest
+
+class RAGForge:
+    """
+    RAGForge Core Orchestrator.
+    Exposes the programmatic API for configuring and running the RAG framework.
+    """
+    def __init__(self, **kwargs):
+        self.config = RAGForgeConfig(**kwargs)
+
+        # Create directories if they do not exist
+        os.makedirs(self.config.RESPONSE_CACHE_DIR, exist_ok=True)
+        os.makedirs(self.config.USER_CACHE_DIR, exist_ok=True)
+        os.makedirs(os.path.dirname(self.config.DATA_PATH) or ".", exist_ok=True)
+
+        # Lazy imports to prevent circular dependency issues
+        from .cache.user_manager import UserManager
+        from .cache.response_cache import ResponseCache
+        from .rag.engine import RAGEngine
+
+        self.user_manager = UserManager(self.config)
+        self.response_cache = ResponseCache(self.config)
+        self.engine = RAGEngine(self.config)
+
+        # Build FastAPI application
+        self.app = create_app(self)
+
+    def run(self, host: str = "127.0.0.1", port: int = 8000, **kwargs):
+        """
+        Starts the FastAPI backend using Uvicorn.
+        """
+        uvicorn.run(self.app, host=host, port=port, **kwargs)
+
+    def ingest(self):
+        """
+        Executes the vector ingestion pipeline for the configured knowledge base text.
+        """
+        ingest(self.config)

ragforge_framework/llm/__init__.py

@@ -0,0 +1 @@
+# LLM Package Initialization

ragforge_framework/llm/gateway.py

@@ -0,0 +1,95 @@
+import contextvars
+import litellm
+from litellm import completion
+
+from ..cache.user_manager import UserManager
+
+# Thread-safe context var to track the client's request IP
+current_ip_var = contextvars.ContextVar("current_ip", default=None)
+
+# Reference to the configuration instance to be read by the global callback
+global_config = None
+
+
+def set_global_config(config):
+    global global_config
+    global_config = config
+
+
+def litellm_success_callback(kwargs, response_obj, start_time, end_time):
+    """
+    Global success callback triggered on every LiteLLM completion.
+    """
+    try:
+        ip = current_ip_var.get()
+        if not ip or not global_config:
+            return
+
+        user_manager = UserManager(global_config)
+        
+        # Get the actual model that completed the request
+        model = response_obj.get("model") or kwargs.get("model") or ""
+        
+        # Detect provider name
+        provider = model.split("/")[0].lower() if "/" in model else "unknown"
+        if provider == "llama3.2:latest":
+            provider = "ollama"
+
+        # Record provider usage
+        user_manager.provider_used(ip, provider)
+
+        # Record fallback if the responding model differs from PRIMARY_MODEL
+        if model != global_config.PRIMARY_MODEL:
+            user_manager.record_fallback(ip)
+    except Exception:
+        pass
+
+
+# Register callback
+litellm.success_callback = [litellm_success_callback]
+
+
+class Gateway:
+    """
+    LiteLLM API Gateway.
+    Handles raw completion calls with auto-fallbacks and retries.
+    """
+    def __init__(self, config):
+        self.config = config
+        self.user_manager = UserManager(config)
+
+        set_global_config(config)
+
+
+    def _detect_provider(self, model: str):
+        return model.split("/")[0].lower()
+
+    def chat(
+        self,
+        ip: str,
+        messages: list,
+        temperature: float = 0.2,
+        timeout: int = 20,
+        retries: int = 3
+    ):
+        # Set the request IP contextvar so the global callback can record stats
+        token = current_ip_var.set(ip)
+        try:
+            response = completion(
+                model=self.config.PRIMARY_MODEL,
+                messages=messages,
+                temperature=temperature,
+                timeout=timeout,
+                num_retries=retries,
+                fallbacks=self.config.FALLBACK_MODELS
+            )
+
+            provider = self._detect_provider(response.model)
+
+            return {
+                "answer": response.choices[0].message.content,
+                "provider": provider,
+                "model": response.model
+            }
+        finally:
+            current_ip_var.reset(token)

ragforge_framework/llm/service.py

@@ -0,0 +1,50 @@
+from langchain_litellm import ChatLiteLLM
+from litellm import completion
+from .gateway import set_global_config
+
+
+class LLMService:
+    """
+    LLM Interaction Service.
+    Wraps native LiteLLM and LangChain ChatLiteLLM model instances.
+    """
+    def __init__(self, config):
+        self.config = config
+        self.model = config.PRIMARY_MODEL
+        self.temperature = config.TEMPERATURE
+        self.timeout = config.TIMEOUT
+        self.max_retries = config.MAX_RETRIES
+        self.fallbacks = config.FALLBACK_MODELS
+
+        set_global_config(config)
+
+    # ===================================================
+    # Native LiteLLM
+    # ===================================================
+
+    def completion(self, messages):
+        return completion(
+            model=self.model,
+            messages=messages,
+            temperature=self.temperature,
+            timeout=self.timeout,
+            num_retries=self.max_retries,
+            fallbacks=self.fallbacks
+        )
+
+    # ===================================================
+    # LangChain Model
+    # ===================================================
+
+    def get_langchain_model(self):
+        model_kwargs = {}
+        if self.fallbacks:
+            model_kwargs["fallbacks"] = self.fallbacks
+
+        return ChatLiteLLM(
+            model=self.model,
+            temperature=self.temperature,
+            timeout=self.timeout,
+            max_retries=self.max_retries,
+            model_kwargs=model_kwargs
+        )

ragforge_framework/models/__init__.py

@@ -0,0 +1 @@
+# Models Package Initialization

ragforge_framework/models/schemas.py

@@ -0,0 +1,90 @@
+from pydantic import BaseModel
+from typing import Dict, List
+
+
+class ChatRequest(BaseModel):
+
+    prompt: str
+
+
+class ChatResponse(BaseModel):
+
+    success: bool
+
+    cached: bool
+
+    answer: str
+
+    remaining_questions: int
+
+    provider: str
+
+
+class UserResponse(BaseModel):
+
+    user_id: str
+
+    ip: str
+
+    total_requests: int
+
+    llm_requests: int
+
+    cache_hits: int
+
+    cache_misses: int
+
+    remaining: int
+
+    first_request: str
+
+    last_request: str
+
+    provider_usage: Dict[str, int]
+
+    fallbacks: int
+
+
+class StatisticsResponse(BaseModel):
+
+    total_users: int
+
+    total_requests: int
+
+    llm_requests: int
+
+    cache_hits: int
+
+    cache_misses: int
+
+    fallbacks: int
+
+    provider_usage: Dict[str, int]
+
+
+# ==========================================================
+# ADMIN RESPONSE SCHEMAS
+# ==========================================================
+
+class AllUsersResponse(BaseModel):
+
+    total_users: int
+
+    users: List[UserResponse]
+
+
+class TopUsersResponse(BaseModel):
+
+    users: List[UserResponse]
+
+
+class CacheResponseDetail(BaseModel):
+
+    answer: str
+
+
+class CacheInformationResponse(BaseModel):
+
+    cached_responses: int
+
+    responses: List[CacheResponseDetail]

ragforge_framework/rag/__init__.py

@@ -0,0 +1 @@
+# RAG Package Initialization