Initial commit (Clean history)

.gitignore

@@ -0,0 +1,20 @@
+# Python-generated files
+__pycache__/
+*.py[oc]
+build/
+dist/
+wheels/
+*.egg-info
+
+# Virtual environments
+.venv
+
+
+#environments
+.env
+
+# Databases
+*.db
+*.sqlite3
+data/chroma_db/
+mlflow.db

.python-version

@@ -0,0 +1 @@
+3.12

Dockerfile

@@ -0,0 +1,27 @@
+# Use official Python runtime
+FROM python:3.10-slim
+
+# Set working directory
+WORKDIR /app
+
+# Install UV
+COPY --from=ghcr.io/astral-sh/uv:latest /uv /bin/uv
+
+# Copy dependency files
+COPY pyproject.toml uv.lock ./
+
+# Install dependencies
+# Note: We exclude mlx here because it's Mac-only. 
+# We install the rest of the project deps.
+RUN uv sync --frozen --no-install-project
+
+# Copy source code
+COPY src ./src
+COPY app ./app
+COPY .env ./.env
+
+# Expose Streamlit port (Must be 7860 for HF Spaces)
+EXPOSE 7860
+
+# Command to run the app
+CMD ["uv", "run", "streamlit", "run", "app/frontend/app.py", "--server.address=0.0.0.0", "--server.port=7860"]

app/frontend/app.py

@@ -0,0 +1,37 @@
+import streamlit as st
+import sys
+import os
+
+# 1. Fix Path
+sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '../../')))
+
+from src.retrieval.rag_chain import build_rag_chain
+
+st.set_page_config(page_title="RAG Observability Platform", layout="wide")
+st.title("🤖 RAG Observability Platform")
+
+if "messages" not in st.session_state:
+    st.session_state.messages = []
+
+@st.cache_resource
+def load_chain():
+    return build_rag_chain()
+
+rag_chain = load_chain()
+
+for message in st.session_state.messages:
+    with st.chat_message(message["role"]):
+        st.markdown(message["content"])
+
+if prompt := st.chat_input("Ask a question..."):
+    st.session_state.messages.append({"role": "user", "content": prompt})
+    with st.chat_message("user"):
+        st.markdown(prompt)
+
+    with st.chat_message("assistant"):
+        with st.spinner("Thinking (M4 GPU)..."):
+            # LCEL Invoke (Direct String)
+            response = rag_chain.invoke(prompt)
+            st.markdown(response)
+            
+    st.session_state.messages.append({"role": "assistant", "content": response})

data/raw/sample.txt

@@ -0,0 +1 @@
+The RAG Observability Platform is a project combining MLX for local inference and Docker for cloud deployment. It uses Dagshub for tracking experiments.

docker/Dockerfile

@@ -0,0 +1,24 @@
+# Use python 3.10 slim image
+FROM python:3.10-slim
+
+# Set working directory
+WORKDIR /app
+
+# Install uv
+COPY --from=ghcr.io/astral-sh/uv:latest /uv /bin/uv
+
+# Copy project files
+COPY pyproject.toml uv.lock ./
+COPY src ./src
+COPY app ./app
+COPY .env ./.env 
+COPY data ./data
+
+# Install dependencies (Excluding the 'local' group which contains MLX)
+RUN uv sync --frozen --no-install-project --no-group local
+
+# Expose the port Streamlit runs on (7860 is mandatory for HF Spaces)
+EXPOSE 7860
+
+# Run the application
+CMD ["uv", "run", "streamlit", "run", "app/frontend/app.py", "--server.port=7860", "--server.address=0.0.0.0"]

main.py

@@ -0,0 +1,6 @@
+def main():
+    print("Hello from rag-observability-platform!")
+
+
+if __name__ == "__main__":
+    main()

pyproject.toml

@@ -0,0 +1,25 @@
+[project]
+name = "rag-observability-platform"
+version = "0.1.0"
+description = "Add your description here"
+readme = "README.md"
+requires-python = ">=3.12"
+dependencies = [
+    "chromadb>=1.3.5",
+    "dagshub>=0.6.3",
+    "langchain>=1.1.2",
+    "langchain-chroma>=1.0.0",
+    "langchain-community>=0.4.1",
+    "langchain-huggingface>=1.1.0",
+    "mlflow>=3.7.0",
+    "pypdf>=6.4.0",
+    "python-dotenv>=1.2.1",
+    "sentence-transformers>=5.1.2",
+    "streamlit>=1.52.1",
+]
+
+[dependency-groups]
+local = [
+    "mlx>=0.30.0",
+    "mlx-lm>=0.28.4",
+]

src/generation/mlx_wrapper.py

@@ -0,0 +1,50 @@
+# src/generation/mlx_wrapper.py
+import os
+from typing import Any, List, Optional
+from langchain_core.callbacks.manager import CallbackManagerForLLMRun
+from langchain_core.language_models.llms import LLM
+from mlx_lm import load, generate
+from dotenv import load_dotenv
+
+load_dotenv()
+
+class MLXLLM(LLM):
+    """Custom LangChain Wrapper for MLX Models"""
+    
+    model_id: str = os.getenv("MODEL_ID", "mlx-community/Llama-3.2-3B-Instruct-4bit")
+    model: Any = None
+    tokenizer: Any = None
+    max_tokens: int = int(os.getenv("MAX_TOKENS", 512))
+
+    def __init__(self, **kwargs):
+        super().__init__(**kwargs)
+        print(f"🚀 Loading MLX Model: {self.model_id}")
+        self.model, self.tokenizer = load(self.model_id)
+
+    @property
+    def _llm_type(self) -> str:
+        return "mlx_llama"
+
+    def _call(
+        self,
+        prompt: str,
+        stop: Optional[List[str]] = None,
+        run_manager: Optional[CallbackManagerForLLMRun] = None,
+        **kwargs: Any,
+    ) -> str:
+        if stop is not None:
+            raise ValueError("stop kwargs are not permitted.")
+
+        messages = [{"role": "user", "content": prompt}]
+        formatted_prompt = self.tokenizer.apply_chat_template(
+            messages, tokenize=False, add_generation_prompt=True
+        )
+
+        response = generate(
+            self.model, 
+            self.tokenizer, 
+            prompt=formatted_prompt, 
+            verbose=False, 
+            max_tokens=self.max_tokens
+        )
+        return response

src/generation/model.py

@@ -0,0 +1,42 @@
+# src/generation/model.py
+import sys
+
+def load_model(model_path="mlx-community/Llama-3.2-3B-Instruct-4bit"):
+    """
+    Loads model conditionally based on environment.
+    Local (Mac): Uses MLX for GPU acceleration.
+    Cloud (Linux): Uses HuggingFace Transformers (CPU/CUDA).
+    """
+    try:
+        from mlx_lm import load, generate
+        print(f"Loading {model_path} with MLX on Apple Silicon...")
+        model, tokenizer = load(model_path)
+        return model, tokenizer, "mlx"
+    except ImportError:
+        # Fallback for Docker/Cloud if MLX isn't available
+        print("MLX not found. Falling back to Transformers...")
+        from transformers import AutoModelForCausalLM, AutoTokenizer
+        tokenizer = AutoTokenizer.from_pretrained(model_path)
+        model = AutoModelForCausalLM.from_pretrained(model_path)
+        return model, tokenizer, "transformers"
+
+if __name__ == "__main__":
+    import mlx.core as mx
+    
+    # 1. Check Default Device
+    device = mx.default_device()
+    print(f"✅ Current MLX Device: {device}")  # Should say "gpu"
+    
+    # 2. Run Inference to trigger GPU
+    model, tokenizer, backend = load_model()
+    
+    if backend == "mlx":
+        from mlx_lm import generate
+        prompt = "Explain quantum physics in one sentence."
+        messages = [{"role": "user", "content": prompt}]
+        prompt_formatted = tokenizer.apply_chat_template(messages, tokenize=False)
+        
+        print(f"\n🧪 Testing Inference (Watch your GPU stats now)...")
+        response = generate(model, tokenizer, prompt=prompt_formatted, verbose=True)
+        print(f"\n🤖 Response: {response}")
+

src/ingestion/pipeline.py

@@ -0,0 +1,81 @@
+# src/ingestion/pipeline.py
+import os
+import mlflow
+import chromadb
+from langchain_community.document_loaders import TextLoader, DirectoryLoader
+from langchain_text_splitters import RecursiveCharacterTextSplitter
+from langchain_huggingface import HuggingFaceEmbeddings
+from langchain_community.vectorstores import Chroma
+from dotenv import load_dotenv
+
+# Load environment variables
+load_dotenv()
+
+# Configuration
+DATA_PATH = "data/raw"
+DB_PATH = "data/chroma_db"
+COLLECTION_NAME = "rag_experiments"
+EMBEDDING_MODEL = "sentence-transformers/all-MiniLM-L6-v2"
+
+class IngestionPipeline:
+    def __init__(self):
+        self.embeddings = HuggingFaceEmbeddings(
+            model_name=EMBEDDING_MODEL,
+            model_kwargs={'device': 'mps'}  # Use M4 MPS for embeddings
+        )
+        
+    def load_documents(self):
+        """Loads text files from the data directory."""
+        loader = DirectoryLoader(DATA_PATH, glob="*.txt", loader_cls=TextLoader)
+        documents = loader.load()
+        print(f"📄 Loaded {len(documents)} documents.")
+        return documents
+
+    def chunk_documents(self, documents, chunk_size=1000, chunk_overlap=200):
+        """Splits documents into smaller chunks."""
+        text_splitter = RecursiveCharacterTextSplitter(
+            chunk_size=chunk_size,
+            chunk_overlap=chunk_overlap
+        )
+        chunks = text_splitter.split_documents(documents)
+        print(f"🧩 Split into {len(chunks)} chunks.")
+        return chunks
+
+    def store_embeddings(self, chunks):
+        """Embeds chunks and stores them in ChromaDB."""
+        if os.path.exists(DB_PATH):
+            print("⚠️ Existing DB found. Appending...")
+            
+        vectorstore = Chroma.from_documents(
+            documents=chunks,
+            embedding=self.embeddings,
+            persist_directory=DB_PATH,
+            collection_name=COLLECTION_NAME
+        )
+        print(f"💾 Saved to {DB_PATH}")
+        return vectorstore
+
+    def run(self):
+        """Runs the full pipeline with MLflow tracking."""
+        mlflow.set_tracking_uri(os.getenv("MLFLOW_TRACKING_URI"))
+        
+        with mlflow.start_run(run_name="Ingestion_Phase_2"):
+            # Log Parameters
+            mlflow.log_param("embedding_model", EMBEDDING_MODEL)
+            mlflow.log_param("chunk_size", 1000)
+            mlflow.log_param("chunk_overlap", 200)
+            
+            # Execute Steps
+            docs = self.load_documents()
+            chunks = self.chunk_documents(docs)
+            self.store_embeddings(chunks)
+            
+            # Log Metrics
+            mlflow.log_metric("num_documents", len(docs))
+            mlflow.log_metric("num_chunks", len(chunks))
+            
+            print("✅ Ingestion complete and logged to Dagshub!")
+
+if __name__ == "__main__":
+    pipeline = IngestionPipeline()
+    pipeline.run()

src/retrieval/rag_chain.py

@@ -0,0 +1,75 @@
+# src/retrieval/rag_chain.py
+import sys
+import os
+from dotenv import load_dotenv
+
+# Add project root to path
+sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '../../')))
+
+import mlflow
+from langchain_chroma import Chroma
+from langchain_huggingface import HuggingFaceEmbeddings
+from langchain_core.prompts import PromptTemplate
+from langchain_core.runnables import RunnablePassthrough
+from langchain_core.output_parsers import StrOutputParser
+from src.generation.mlx_wrapper import MLXLLM
+
+load_dotenv()
+
+# Configuration from ENV
+DB_PATH = os.getenv("CHROMA_DB_PATH", "data/chroma_db")
+COLLECTION_NAME = os.getenv("CHROMA_COLLECTION_NAME", "rag_experiments")
+EMBEDDING_MODEL = os.getenv("EMBEDDING_MODEL", "sentence-transformers/all-MiniLM-L6-v2")
+
+def format_docs(docs):
+    return "\n\n".join(doc.page_content for doc in docs)
+
+def build_rag_chain():
+    """Builds and returns the RAG chain using LCEL."""
+    
+    # 1. Initialize Embeddings
+    embeddings = HuggingFaceEmbeddings(
+        model_name=EMBEDDING_MODEL,
+        model_kwargs={'device': 'mps'}
+    )
+    
+    # 2. Initialize Retriever
+    vectorstore = Chroma(
+        persist_directory=DB_PATH,
+        embedding_function=embeddings,
+        collection_name=COLLECTION_NAME
+    )
+    retriever = vectorstore.as_retriever(search_kwargs={"k": 2})
+    
+    # 3. Initialize LLM (No arguments needed, it pulls from env)
+    llm = MLXLLM()
+    
+    # 4. Create Prompt Template
+    template = """Use the following pieces of context to answer the question at the end.
+    If you don't know the answer, just say that you don't know, don't try to make up an answer.
+    
+    Context: {context}
+    
+    Question: {question}
+    
+    Answer:"""
+    custom_prompt = PromptTemplate.from_template(template)
+    
+    # 5. Build LCEL Chain
+    rag_chain = (
+        {"context": retriever | format_docs, "question": RunnablePassthrough()}
+        | custom_prompt
+        | llm
+        | StrOutputParser()
+    )
+    
+    return rag_chain
+
+if __name__ == "__main__":
+    chain = build_rag_chain()
+    query = "What technologies does the RAG platform use?"
+    
+    print(f"\n❓ Query: {query}")
+    res = chain.invoke(query)
+    print("\n🤖 Answer:")
+    print(res)