Sync from GitHub

.dockerignore

@@ -0,0 +1,31 @@
+# Docker ignore file - exclude unnecessary files from Docker context
+__pycache__/
+*.pyc
+*.pyo
+*.pyd
+.Python
+.pytest_cache/
+.coverage
+htmlcov/
+*.egg-info/
+dist/
+build/
+.git/
+.gitignore
+.env
+.env.local
+.vscode/
+.idea/
+*.log
+*.db
+.DS_Store
+node_modules/
+.venv/
+venv/
+env/
+*.swp
+*.swo
+*~
+.github/
+README.md
+TEST_SUMMARY.md

.gitattributes

@@ -1,35 +0,0 @@
-*.7z filter=lfs diff=lfs merge=lfs -text
-*.arrow filter=lfs diff=lfs merge=lfs -text
-*.bin filter=lfs diff=lfs merge=lfs -text
-*.bz2 filter=lfs diff=lfs merge=lfs -text
-*.ckpt filter=lfs diff=lfs merge=lfs -text
-*.ftz filter=lfs diff=lfs merge=lfs -text
-*.gz filter=lfs diff=lfs merge=lfs -text
-*.h5 filter=lfs diff=lfs merge=lfs -text
-*.joblib filter=lfs diff=lfs merge=lfs -text
-*.lfs.* filter=lfs diff=lfs merge=lfs -text
-*.mlmodel filter=lfs diff=lfs merge=lfs -text
-*.model filter=lfs diff=lfs merge=lfs -text
-*.msgpack filter=lfs diff=lfs merge=lfs -text
-*.npy filter=lfs diff=lfs merge=lfs -text
-*.npz filter=lfs diff=lfs merge=lfs -text
-*.onnx filter=lfs diff=lfs merge=lfs -text
-*.ot filter=lfs diff=lfs merge=lfs -text
-*.parquet filter=lfs diff=lfs merge=lfs -text
-*.pb filter=lfs diff=lfs merge=lfs -text
-*.pickle filter=lfs diff=lfs merge=lfs -text
-*.pkl filter=lfs diff=lfs merge=lfs -text
-*.pt filter=lfs diff=lfs merge=lfs -text
-*.pth filter=lfs diff=lfs merge=lfs -text
-*.rar filter=lfs diff=lfs merge=lfs -text
-*.safetensors filter=lfs diff=lfs merge=lfs -text
-saved_model/**/* filter=lfs diff=lfs merge=lfs -text
-*.tar.* filter=lfs diff=lfs merge=lfs -text
-*.tar filter=lfs diff=lfs merge=lfs -text
-*.tflite filter=lfs diff=lfs merge=lfs -text
-*.tgz filter=lfs diff=lfs merge=lfs -text
-*.wasm filter=lfs diff=lfs merge=lfs -text
-*.xz filter=lfs diff=lfs merge=lfs -text
-*.zip filter=lfs diff=lfs merge=lfs -text
-*.zst filter=lfs diff=lfs merge=lfs -text
-*tfevents* filter=lfs diff=lfs merge=lfs -text

.github/workflows/hugging_face_sync.yml

@@ -0,0 +1,32 @@
+name: Sync to Hugging Face Space
+
+on:
+  push:
+    branches:
+      - main
+
+jobs:
+  sync:
+    runs-on: ubuntu-latest
+
+    steps:
+      - name: Checkout GitHub repo
+        uses: actions/checkout@v3
+
+      - name: Clone Hugging Face Space
+        run: |
+          git clone https://mukiibi:${{ secrets.HF_TOKEN_RON }}@huggingface.co/spaces/Sebunya/AskXeno hf_space
+
+      - name: Sync files
+        run: |
+          rsync -av --delete --exclude ".git" ./ hf_space/ --ignore-errors --exclude "hf_space"
+
+
+      - name: Commit & push
+        run: |
+          cd hf_space
+          git config user.email "actions@github.com"
+          git config user.name "github-actions"
+          git add .
+          git commit -m "Sync from GitHub" || echo "No changes to commit"
+          git push origin main

.github/workflows/tests.yml

@@ -0,0 +1,48 @@
+name: Run Tests
+
+on:
+  push:
+    branches: [ main, develop ]
+  pull_request:
+    branches: [ main, develop ]
+
+jobs:
+  test:
+    runs-on: ubuntu-latest
+    
+    strategy:
+      matrix:
+        python-version: ['3.11']
+    
+    steps:
+    - uses: actions/checkout@v3
+    
+    - name: Set up Python ${{ matrix.python-version }}
+      uses: actions/setup-python@v4
+      with:
+        python-version: ${{ matrix.python-version }}
+    
+    - name: Cache dependencies
+      uses: actions/cache@v3
+      with:
+        path: ~/.cache/pip
+        key: ${{ runner.os }}-pip-${{ hashFiles('**/requirements.txt') }}
+        restore-keys: |
+          ${{ runner.os }}-pip-
+    
+    - name: Install dependencies
+      run: |
+        python -m pip install --upgrade pip
+        pip install -r requirements.txt
+    
+    - name: Run tests with coverage
+      run: |
+        pytest --cov=src --cov-report=xml --cov-report=term-missing
+    
+    - name: Upload coverage to Codecov
+      uses: codecov/codecov-action@v3
+      with:
+        file: ./coverage.xml
+        flags: unittests
+        name: codecov-umbrella
+        fail_ci_if_error: false

.gitignore

@@ -0,0 +1,30 @@
+# Ignore environment variable files
+.env
+.venv/
+
+# Ignore application log files
+app.log
+
+# Ignore compiled Python files
+__pycache__/
+
+# Ignore test coverage reports
+htmlcov/
+.coverage
+coverage.xml
+.pytest_cache/
+
+# Ignore test artifacts
+.tox/
+.cache/
+*.pyc
+*.pyo
+*.pyc
+
+# Ignore database files
+*.db
+
+# Ignore binary files
+*.txt
+*.bin
+*.exe

DOCKER_BUILD_INFO.md

@@ -0,0 +1,94 @@
+# Docker Build Instructions 
+
+## Dockerfile Optimizations
+
+```dockerfile
+ENV PIP_DEFAULT_TIMEOUT=100 \
+    PIP_RETRIES=5
+
+RUN pip install --upgrade pip setuptools wheel && \
+    pip install --no-cache-dir --default-timeout=100 -r requirements.txt
+```
+
+## Build Progress
+
+The Docker build is currently in progress. Here's what it's doing:
+
+1. ✅ Loading base Python 3.10-slim image
+2. ✅ Setting working directory and environment variables
+3. ⏳ Installing system dependencies (build-essential, curl)
+4. ⏳ Installing Python dependencies (this may take 10-15 minutes due to large packages like PyTorch ~900MB)
+5. Copying project files
+6. Creating necessary directories
+7. Exposing port 7860
+8. Setting health check
+
+## How to Use
+
+### Start the Docker Container
+
+Once the build completes, start with:
+
+```bash
+docker-compose up -d
+```
+
+### Access the Application
+
+- Open browser to `http://localhost:7860`
+
+### View Logs
+
+```bash
+docker-compose logs -f xeno-bot
+```
+
+### Stop the Container
+
+```bash
+docker-compose down
+```
+
+### Build Directly Without docker-compose
+
+```bash
+docker build -t xeno-bot:latest .
+docker run -p 7860:7860 \
+  -e GEMINI_API_KEY="your-api-key" \
+  -e GOOGLE_SHEETS_CREDENTIALS='{"type": "service_account", ...}' \
+  xeno-bot:latest
+```
+
+## Environment Variables Required
+
+The container needs these environment variables set in `.env`:
+
+```
+GEMINI_API_KEY=your-google-gemini-api-key
+GOOGLE_SHEETS_CREDENTIALS={"type": "service_account", "project_id": "...", ...}
+```
+
+See `.env.example` for template.
+
+## Performance Notes
+
+- **Initial build time**: 10-20 minutes (downloading and installing ~900MB PyTorch library)
+- **Subsequent builds**: Faster due to Docker layer caching
+- **Container startup**: ~30-60 seconds for first run, ~5-10 seconds after that
+- **Memory requirement**: 2GB minimum recommended (PyTorch + Gradio + ChromaDB)
+
+## Troubleshooting
+
+If the build still times out:
+1. Increase `PIP_DEFAULT_TIMEOUT` further in Dockerfile
+2. Check your network connection
+3. Try building again (Docker will use cached layers)
+4. Consider using a build cache: `DOCKER_BUILDKIT=1 docker build ...`
+
+## Files Modified
+
+- `Dockerfile` - Optimized for reliability
+- `docker-compose.yml` - Removed obsolete version attribute
+- `app.py` - Added missing `import os`
+- `.dockerignore` - Excludes unnecessary files from build context
+- `.env.example` - Template for environment variables

Dockerfile

@@ -0,0 +1,47 @@
+# Use official Python runtime as base image
+FROM python:3.11-slim
+
+# Set working directory
+WORKDIR /app
+
+# Set environment variables
+ENV PYTHONUNBUFFERED=1 \
+    PYTHONDONTWRITEBYTECODE=1 \
+    GRADIO_SERVER_NAME="0.0.0.0" \
+    GRADIO_SERVER_PORT=7860 \
+    PIP_DEFAULT_TIMEOUT=100 \
+    PIP_RETRIES=5
+
+# Install system dependencies
+RUN apt-get update && apt-get install -y --no-install-recommends \
+    build-essential \
+    curl \
+    && rm -rf /var/lib/apt/lists/*
+
+# Copy requirements first for better layer caching
+COPY requirements.txt .
+
+# Upgrade pip and install dependencies with retries
+RUN pip install --upgrade pip setuptools wheel && \
+    pip install --no-cache-dir --default-timeout=100 --prefer-binary \
+        torch==2.3.1+cpu --index-url https://download.pytorch.org/whl/cpu && \
+    pip install --no-cache-dir --default-timeout=100 --prefer-binary \
+        -r requirements.txt --extra-index-url https://download.pytorch.org/whl/cpu
+
+# Copy project files
+COPY app.py .
+COPY src/ ./src/
+COPY XENO_Uganda_KnowledgeBase_Advisory.json ./
+
+# Create necessary directories
+RUN mkdir -p /tmp/xeno_db
+
+# Expose Gradio port
+EXPOSE 7860
+
+# Health check
+HEALTHCHECK --interval=30s --timeout=10s --start-period=5s --retries=3 \
+    CMD curl -f http://localhost:7860 || exit 1
+
+# Run the application
+CMD ["python", "app.py"]

README.md

@@ -9,39 +9,52 @@ app_file: app.py
 pinned: false
 ---
 
-ASKXENO - AI-Powered XENO Support Assistant
-Overview
+# ASKXENO - AI-Powered XENO Support Assistant
+## Overview
 ASKXENO is an AI-powered customer support assistant designed to provide accurate and timely responses to queries about XENO financial services. Built with a Retrieval-Augmented Generation (RAG) pipeline, it leverages a knowledge base, intent classification, and conversation memory to deliver professional and context-aware responses. The application is deployed as a Hugging Face Space and includes performance tracking for response time analysis.
 Features
 
 Natural Language Query Handling: Responds to user questions about XENO services, including account management, transactions, platform features, and general information.
+
 Intent Classification: Detects simple intents (e.g., greetings, thanks, goodbyes) for quick, tailored responses without querying the knowledge base.
 RAG Pipeline: Uses a ChromaDB vector store and Google Gemini embeddings to retrieve relevant information from a JSON-based XENO knowledge base.
 Conversation Memory: Maintains chat history using LangGraph's SqliteSaver for context-aware responses.
+
 Performance Tracking: Logs response times for each processing step (e.g., intent classification, retrieval, LLM generation) to a Google Sheet for analysis.
+
 Gradio UI: Provides an interactive, user-friendly interface with session tracking and a chatbot-style conversation display.
+
 Error Handling: Includes robust logging to handle and record errors, with fallback local file logging if Google Sheets fails.
 
-Tech Stack
+## Tech Stack
+
+**Python Libraries:**
 
-Python Libraries:
 gradio: For the web-based user interface.
+
 pandas: For handling the JSON knowledge base.
+
 sentence_transformers: For similarity calculations.
+
 google.generativeai: For embeddings and LLM responses (Gemini API).
+
 chromadb: For vector storage and retrieval.
+
 langchain_chroma: For integrating ChromaDB with LangChain.
+
 gspread: For logging to Google Sheets.
+
 langgraph: For conversation memory management.
+
 torch: For tensor operations in similarity calculations.
 
 
-Database:
-SQLite for conversation memory (xeno_memory.db).
+## Database:
+SQLite for conversation memory (xeno_memory.db). [currently iys implemented in google sheets]
 ChromaDB for persistent vector storage (/tmp/xeno_db).
 
 
-External Services:
+## External Services:
 Google Sheets for logging responses and timing data.
 Google Gemini API for embeddings and text generation.
 
@@ -50,10 +63,6 @@ Google Gemini API for embeddings and text generation.
 Setup and Installation
 This project is designed to run in a Hugging Face Space. To set it up locally or in a similar environment, follow these steps:
 
-Clone the Repository (if not running in Hugging Face Space):
-git clone <repository-url>
-cd <repository-directory>
-
 
 Install Dependencies:Ensure Python 3.8+ is installed, then install the required packages:
 pip install -r requirements.txt

TEST_SUMMARY.md

@@ -0,0 +1,208 @@
+# XENO Bot - Test Suite Summary
+
+## Overview
+
+Comprehensive unit test suite created for the XENO Bot application with **67 passing tests** and **88% code coverage**.
+
+## Test Statistics
+
+- **Total Tests**: 67
+- **All Passing**: ✅ 100%
+- **Overall Coverage**: 88%
+- **Test Files**: 7
+- **Modules Tested**: 7
+
+## Module Coverage
+
+| Module | Statements | Tested | Coverage | Missing Lines |
+|--------|------------|--------|----------|---------------|
+| config.py | 22 | 21 | 95% | 11 |
+| intent_classifier.py | 24 | 24 | **100%** | - |
+| knowledge_base.py | 20 | 20 | **100%** | - |
+| logger.py | 58 | 52 | 90% | 28, 56-66, 92 |
+| memory.py | 30 | 28 | 93% | 31, 69 |
+| response_generator.py | 23 | 22 | 96% | 27 |
+| utils.py | 35 | 32 | 91% | 21-23 |
+| vector_store.py | 65 | 46 | 71% | 30-57, 76, 117, 148 |
+
+## Test Files
+
+### 1. test_utils.py (8 tests)
+Tests the `PipelineTimer` class for timing pipeline execution.
+
+**Tests:**
+- Timer initialization
+- Reset functionality
+- Context manager timing
+- Multiple step timing
+- Total time calculation
+- Timing summary generation
+- Current step tracking
+- Exception handling
+
+### 2. test_intent_classifier.py (12 tests)
+Tests the `IntentClassifier` class for user intent classification.
+
+**Tests:**
+- Classification of greetings, thanks, goodbye, and queries
+- Case insensitivity
+- Timer integration
+- Simple intent detection
+- Dynamic intent addition
+- Response variety
+- Empty and mixed messages
+
+### 3. test_knowledge_base.py (8 tests)
+Tests knowledge base loading and document preparation.
+
+**Tests:**
+- JSON file loading
+- Null content filtering
+- Document preparation
+- Metadata structure
+- Missing field handling
+- Empty knowledge base
+- Document text formatting
+
+### 4. test_memory.py (9 tests)
+Tests LangGraph memory operations with SQLite.
+
+**Tests:**
+- Session config creation
+- Memory update and retrieval
+- Empty checkpoint handling
+- Timer integration
+- Checkpoint structure validation
+
+### 5. test_response_generator.py (10 tests)
+Tests LLM response generation functionality.
+
+**Tests:**
+- Chat history formatting
+- Response generation
+- Prompt structure
+- System prompt inclusion
+- Timer integration
+- Empty history handling
+- Text stripping
+
+### 6. test_logger.py (10 tests)
+Tests Google Sheets logging functionality.
+
+**Tests:**
+- Response logging
+- Timing data logging
+- Error handling and fallback
+- Empty/single knowledge pairs
+- Long question truncation
+- Missing step times
+
+### 7. test_vector_store.py (10 tests)
+Tests ChromaDB vector store operations.
+
+**Tests:**
+- Embedding generation
+- Similarity calculation
+- Context processing
+- Multiple documents
+- Result limiting
+- Missing metadata handling
+- Timer integration
+
+## Running Tests
+
+### Quick Start
+```bash
+# Install test dependencies
+pip install pytest pytest-cov pytest-mock
+
+# Run all tests
+pytest
+
+# Run with coverage
+pytest --cov=src --cov-report=html
+
+# Run specific test file
+pytest tests/test_utils.py
+
+# Run with verbose output
+pytest -v
+```
+
+### Using the Test Runner
+```bash
+python run_tests.py
+```
+
+## CI/CD Integration
+
+GitHub Actions workflow created at `.github/workflows/tests.yml`
+
+**Features:**
+- Runs on push and pull requests
+- Tests against Python 3.9, 3.10, and 3.11
+- Generates coverage reports
+- Uploads to Codecov
+
+## Test Configuration
+
+### pytest.ini
+- Defines test discovery patterns
+- Sets test collection rules
+
+### setup.cfg
+- Coverage configuration
+- Exclude patterns for coverage
+- Report formatting
+
+### conftest.py
+- Shared fixtures and mocks
+- Environment setup
+- Mock external services (Google Sheets, Google AI, ChromaDB)
+
+## Best Practices Implemented
+
+1. **Isolation**: Tests don't depend on external services
+2. **Mocking**: External APIs and databases are mocked
+3. **Coverage**: High coverage across all critical modules
+4. **Documentation**: Clear test descriptions and docstrings
+5. **Structure**: Consistent test organization
+6. **CI/CD Ready**: Automated testing pipeline configured
+
+## Mock Strategy
+
+To avoid dependencies on external services during testing:
+
+- **Google Generative AI**: Mocked at import time
+- **Google Sheets**: Mocked gspread and oauth2 modules
+- **ChromaDB**: Mocked PersistentClient
+- **SQLite**: Mocked connections for memory tests
+
+## Future Improvements
+
+1. Increase vector_store.py coverage (currently 71%)
+2. Add integration tests for full pipeline
+3. Add performance benchmarking tests
+4. Add tests for Gradio interface (app.py)
+5. Add stress tests for concurrent requests
+
+## Contributing
+
+When adding new features:
+
+1. Write tests first (TDD approach)
+2. Ensure tests pass locally
+3. Check coverage doesn't drop below 85%
+4. Update this documentation
+
+## Test Execution Time
+
+- **Average runtime**: ~8-17 seconds
+- **Fastest**: ~2 seconds (utils + intent_classifier only)
+- **With coverage**: ~16-17 seconds
+
+---
+
+**Generated**: December 10, 2025  
+**Test Framework**: pytest 9.0.2  
+**Python Version**: 3.13.9

app.py

@@ -1,5 +1,9 @@
-import uuid
+"""
+XENO Bot - AI-powered customer service assistant
+Main application file with Gradio interface
+"""
 import os
+import uuid
 import gradio as gr
 import pandas as pd
 import torch
@@ -19,64 +23,24 @@ from typing import Dict, List, Tuple
 import time
 from contextlib import contextmanager
 import threading  # <--- Added for non-blocking feedback logging
-
 import logging
 import traceback
-import sys
 
-# ===== Configure Logging =====
-logging.basicConfig(
-    filename="app.log",
-    level=logging.INFO,
-    format="%(asctime)s - %(levelname)s - %(message)s"
+# Import custom modules
+from src.utils import PipelineTimer
+from src.config import SIMILARITY_THRESHOLD, SERVER_NAME, SERVER_PORT
+from src.memory import create_session_config, update_memory, retrieve_memory
+from src.intent_classifier import IntentClassifier
+from src.vector_store import (
+    initialize_vector_store,
+    generate_embeddings,
+    calculate_similarity,
+    process_context
 )
+from src.response_generator import generate_xeno_response
+from src.logger import log_response, log_timing_data
 
-def log_exception(exc_type, exc_value, exc_traceback):
-    if issubclass(exc_type, KeyboardInterrupt):
-        return
-    logging.critical("Uncaught exception", exc_info=(exc_type, exc_value, exc_traceback))
-
-sys.excepthook = log_exception
-logging.info("App started successfully.")
-
-# ===== Time Tracking Class =====
-class PipelineTimer:
-    def __init__(self):
-        self.reset()
-    
-    def reset(self):
-        """Reset all timing data for a new request"""
-        self.start_time = time.time()
-        self.step_times = {}
-        self.step_start = None
-        self.current_step = None
-    
-    @contextmanager
-    def time_step(self, step_name: str):
-        """Context manager to time a specific step"""
-        step_start = time.time()
-        self.current_step = step_name
-        try:
-            yield
-        finally:
-            step_end = time.time()
-            self.step_times[step_name] = round((step_end - step_start) * 1000, 2)  # Convert to milliseconds
-            self.current_step = None
-    
-    def get_total_time(self):
-        """Get total elapsed time since reset"""
-        return round((time.time() - self.start_time) * 1000, 2)
-    
-    def get_timing_summary(self):
-        """Get a summary of all timing data"""
-        total_time = self.get_total_time()
-        return {
-            'total_time_ms': total_time,
-            'step_times': self.step_times,
-            'timestamp': datetime.now().isoformat()
-        }
-
-# Initialize global timer
+# Initialize components
 timer = PipelineTimer()
 
 # === Configuration ===
@@ -410,11 +374,11 @@ def get_context_and_answer(message, history, session_id="default"):
     notes = []
     
     try:
-        config = {"configurable": {"thread_id": str(session_id), "checkpoint_ns": ""}}
+        # Create session config
+        config = create_session_config(session_id)
         
         # Step 1: Intent Classification
-        with timer.time_step("intent_classification"):
-            intent, direct_response = intent_classifier.classify_intent(message)
+        intent, direct_response = intent_classifier.classify_intent(message)
         
         # Step 2: Memory Retrieval
         chat_history = retrieve_memory(config)
@@ -437,21 +401,9 @@ def get_context_and_answer(message, history, session_id="default"):
                         queried_results = retriever.invoke(message)
                     
                     # Step 4: Embedding Generation
-                    with timer.time_step("embedding_generation"):
-                        query_embedding = genai.embed_content(
-                            model=embedding_model, 
-                            content=message, 
-                            task_type="retrieval_query"
-                        )['embedding']
-                        
-                        doc_embeddings = [
-                            genai.embed_content(
-                                model=embedding_model, 
-                                content=doc.page_content, 
-                                task_type="retrieval_document"
-                            )['embedding'] 
-                            for doc in queried_results
-                        ]
+                    query_embedding, doc_embeddings = generate_embeddings(
+                        message, queried_results, timer
+                    )
                     
                     # Step 5: Similarity Calculation
                     with timer.time_step("similarity_calculation"):
@@ -461,7 +413,7 @@ def get_context_and_answer(message, history, session_id="default"):
                         )[0].tolist()
                         max_score = max(cosine_scores) if cosine_scores else 0
 
-                    if max_score < 0.4:
+                    if max_score < SIMILARITY_THRESHOLD:
                         answer = "I'm sorry, I couldn't find specific information for your question. Could you try rephrasing it, or contact XENO support directly?"
                         notes.append(f"Low similarity score: {max_score:.3f}")
                     else:
@@ -484,8 +436,7 @@ def get_context_and_answer(message, history, session_id="default"):
         update_memory(config, message, answer)
         
         # Step 9: Response Logging
-        with timer.time_step("response_logging"):
-            log_response(message, answer, source_ids, knowledge_pairs, session_id)
+        log_response(message, answer, source_ids, knowledge_pairs, session_id)
         
         # Log timing data
         timing_summary = timer.get_timing_summary()
@@ -515,8 +466,9 @@ def get_context_and_answer(message, history, session_id="default"):
         
         return "I apologize, but I encountered an error processing your request. Please try again."
 
+
 # === Enhanced Gradio UI ===
-def respond(message, history, session_id):
+def respond(message: str, history: List, session_id: str):
     """Gradio's main response function"""
     if not session_id:
         session_id = str(uuid.uuid4())
@@ -526,7 +478,9 @@ def respond(message, history, session_id):
     
     return "", history
 
+
 def create_interface():
+    """Create Gradio interface"""
     with gr.Blocks(theme=gr.themes.Soft()) as demo:
         gr.Markdown("""
         # ASKXENO
@@ -593,6 +547,12 @@ def create_interface():
             
     return demo
 
+
 if __name__ == "__main__":
     iface = create_interface()
-    iface.launch(share=False, server_name="0.0.0.0", server_port=7860, ssr_mode=False)
+    iface.launch(
+        share=False, 
+        server_name=SERVER_NAME, 
+        server_port=SERVER_PORT, 
+        ssr_mode=False
+    )

docker-compose.yml

@@ -0,0 +1,27 @@
+services:
+  xeno-bot:
+    build:
+      context: .
+      dockerfile: Dockerfile
+    ports:
+      - "7860:7860"
+    environment:
+      - GEMINI_API_KEY=${GEMINI_API_KEY}
+      - GOOGLE_SHEETS_CREDENTIALS=${GOOGLE_SHEETS_CREDENTIALS}
+    volumes:
+      - xeno_db:/tmp/xeno_db
+      - chroma_cache:/root/.cache/chroma
+      - ./xeno_memory.db:/app/xeno_memory.db
+    restart: unless-stopped
+    healthcheck:
+      test: ["CMD", "python", "-c", "import sys, urllib.request; sys.exit(0) if urllib.request.urlopen('http://localhost:7860').status == 200 else sys.exit(1)"]
+      interval: 30s
+      timeout: 10s
+      retries: 3
+      start_period: 5s
+
+volumes:
+  xeno_db:
+    driver: local
+  chroma_cache:
+    driver: local

newnewfile.py

@@ -1 +0,0 @@
-print('hello')

pytest.ini

@@ -0,0 +1,5 @@
+[pytest]
+testpaths = tests
+python_files = test_*.py
+python_classes = Test*
+python_functions = test_*

requirements.txt

@@ -2,7 +2,7 @@ huggingface_hub==0.25.2
 gradio
 pydantic==2.10.6
 pandas
-torch
+torch==2.3.1+cpu
 numpy
 sentence-transformers
 google-generativeai
@@ -13,4 +13,9 @@ langchain-chroma
 gspread
 google-auth
 
-python-dateutil
+python-dateutil
+
+# Testing dependencies
+pytest>=7.0.0
+pytest-cov>=4.0.0
+pytest-mock>=3.10.0

run_tests.py

@@ -0,0 +1,38 @@
+"""
+Test runner script for XENO Bot
+Run this script to execute all tests with coverage reporting
+"""
+import sys
+import subprocess
+
+
+def run_tests():
+    """Run all tests with coverage"""
+    print("=" * 70)
+    print("Running XENO Bot Unit Tests")
+    print("=" * 70)
+    
+    # Run pytest with coverage
+    cmd = [
+        sys.executable, "-m", "pytest",
+        "--cov=src",
+        "--cov-report=term-missing",
+        "--cov-report=html",
+        "-v"
+    ]
+    
+    result = subprocess.run(cmd)
+    
+    print("\n" + "=" * 70)
+    if result.returncode == 0:
+        print("✓ All tests passed!")
+        print("Coverage report generated in htmlcov/index.html")
+    else:
+        print("✗ Some tests failed!")
+    print("=" * 70)
+    
+    return result.returncode
+
+
+if __name__ == "__main__":
+    sys.exit(run_tests())

setup.cfg

@@ -0,0 +1,28 @@
+[tool:pytest]
+testpaths = tests
+python_files = test_*.py
+python_classes = Test*
+python_functions = test_*
+addopts = 
+    -v
+    --tb=short
+    --strict-markers
+    --disable-warnings
+
+# Coverage options
+[coverage:run]
+source = src
+omit = 
+    */tests/*
+    */test_*
+    */__pycache__/*
+
+[coverage:report]
+exclude_lines =
+    pragma: no cover
+    def __repr__
+    raise AssertionError
+    raise NotImplementedError
+    if __name__ == .__main__.:
+    if TYPE_CHECKING:
+    @abstractmethod

src/__init__.py

@@ -0,0 +1,3 @@
+"""
+XENO Bot source package
+"""

src/config.py

@@ -0,0 +1,48 @@
+"""
+Configuration module for XENO Bot
+Handles environment variables and application settings
+"""
+import os
+import google.generativeai as genai
+
+# === API Configuration ===
+GEMINI_API_KEY = os.environ.get("GEMINI_API_KEY")
+if not GEMINI_API_KEY:
+    raise ValueError("GEMINI_API_KEY environment variable not set.")
+
+genai.configure(api_key=GEMINI_API_KEY)
+
+# === Model Configuration ===
+EMBEDDING_MODEL = "models/embedding-001"
+LLM_MODEL_NAME = "models/gemma-3-4b-it"
+
+# === Database Configuration ===
+COLLECTION_NAME = "xeno_collection"
+CHROMA_DB_PATH = "/tmp/xeno_db"
+SQLITE_DB_PATH = "xeno_memory.db"
+
+# === Knowledge Base Configuration ===
+KNOWLEDGE_BASE_PATH = "XENO_Uganda_KnowledgeBase_Advisory.json"
+
+# === Google Sheets Configuration ===
+GOOGLE_SHEETS_CREDENTIALS_ENV = "GOOGLE_SHEETS_CREDENTIALS"
+SPREADSHEET_NAME = "Response_Log"
+RESPONSE_SHEET_INDEX = 0  # sheet1
+TIMING_SHEET_NAME = "Timing_Log"
+
+# === RAG Configuration ===
+RAG_TOP_K = 4
+RAG_MAX_RESULTS = 2
+SIMILARITY_THRESHOLD = 0.4
+
+# === Server Configuration ===
+SERVER_NAME = "0.0.0.0"
+SERVER_PORT = 7860
+
+# === Prompt Configuration ===
+SYSTEM_PROMPT = """You are a friendly XENO Support Assistant, an AI-powered helpful and professional customer service representative.
+Use only the information provided in the knowledge base context to answer user queries.
+Do not hallucinate. If context doesn't contain relevant info, say so in a calm polite manner by saying I'm sorry, I can't assist with that.
+Only use context that is clearly relevant to the user's question.
+For greetings like "hi" or "hello", respond politely without using the context.
+remember previous conversations."""

src/intent_classifier.py

@@ -0,0 +1,107 @@
+"""
+Intent Classification module for XENO Bot
+Handles classification of user intents (greetings, thanks, goodbye, queries)
+"""
+import re
+import random
+from typing import Tuple, List
+
+
+class IntentClassifier:
+    """Classifies user intents and provides appropriate responses"""
+    
+    def __init__(self):
+        self.intent_patterns = {
+            'greeting': {
+                'patterns': [
+                    r'\b(hi|hello|hey|good morning|good afternoon|good evening|greetings)\b',
+                    r'^(hi|hello|hey)[\s!.]*$',
+                    r'\b(how are you|how do you do)\b'
+                ],
+                'responses': [
+                    "Hello! I'm XENO Assistant. How can I help you with XENO financial services today?",
+                    "Hi there! I'm here to assist you with any questions about XENO services. What can I help you with?",
+                    "Good day! Welcome to XENO Support. How may I assist you today?"
+                ]
+            },
+            'thanks': {
+                'patterns': [
+                    r'\b(thank you|thanks|thank u|thx|appreciate|grateful)\b',
+                    r'^(thanks|thank you)[\s!.]*$',
+                    r'\b(much appreciated|thanks a lot|thank you so much)\b'
+                ],
+                'responses': [
+                    "You're welcome! Is there anything else I can help you with regarding XENO services?",
+                    "Happy to help! Feel free to ask if you have any other questions about XENO.",
+                    "Glad I could assist you! Let me know if you need help with anything else."
+                ]
+            },
+            'goodbye': {
+                'patterns': [
+                    r'\b(bye|goodbye|see you|farewell|take care|have a good day)\b',
+                    r'^(bye|goodbye)[\s!.]*$',
+                    r'\b(talk to you later|see you later|until next time)\b'
+                ],
+                'responses': [
+                    "Goodbye! Thank you for using XENO services. Have a great day!",
+                    "Take care! Feel free to return anytime you need help with XENO services.",
+                    "Have a wonderful day! Don't hesitate to reach out if you need assistance with XENO."
+                ]
+            }
+        }
+    
+    def classify_intent(self, message: str, timer=None) -> Tuple[str, str]:
+        """
+        Classify the intent of a user message
+        
+        Args:
+            message: User's message
+            timer: Optional timer object for tracking
+        
+        Returns:
+            Tuple of (intent_name, response_text)
+        """
+        if timer:
+            with timer.time_step("intent_classification"):
+                return self._classify_intent_impl(message)
+        else:
+            return self._classify_intent_impl(message)
+    
+    def _classify_intent_impl(self, message: str) -> Tuple[str, str]:
+        """Internal implementation of intent classification"""
+        message_lower = message.lower().strip()
+        
+        for intent_name, intent_data in self.intent_patterns.items():
+            for pattern in intent_data['patterns']:
+                if re.search(pattern, message_lower, re.IGNORECASE):
+                    response = random.choice(intent_data['responses'])
+                    return intent_name, response
+        
+        return 'query', ''
+    
+    def is_simple_intent(self, intent: str) -> bool:
+        """
+        Check if the intent is a simple one that doesn't require RAG
+        
+        Args:
+            intent: Intent name
+        
+        Returns:
+            True if simple intent, False otherwise
+        """
+        simple_intents = ['greeting', 'thanks']
+        return intent in simple_intents
+    
+    def add_intent(self, intent_name: str, patterns: List[str], responses: List[str]):
+        """
+        Add a new intent to the classifier
+        
+        Args:
+            intent_name: Name of the intent
+            patterns: List of regex patterns to match
+            responses: List of possible responses
+        """
+        self.intent_patterns[intent_name] = {
+            'patterns': patterns,
+            'responses': responses
+        }

src/knowledge_base.py

@@ -0,0 +1,69 @@
+"""
+Knowledge Base module for XENO Bot
+Handles loading and preparing knowledge base data
+"""
+import pandas as pd
+from typing import List, Dict, Tuple, Any
+from src.config import KNOWLEDGE_BASE_PATH
+
+
+def load_knowledge_base(filepath: str = KNOWLEDGE_BASE_PATH) -> pd.DataFrame:
+    """
+    Load knowledge base from JSON file
+    
+    Args:
+        filepath: Path to the knowledge base JSON file
+    
+    Returns:
+        DataFrame with knowledge base data
+    """
+    df = pd.read_json(filepath)
+    df.dropna(subset=['Content'], inplace=True)
+    return df
+
+
+def prepare_documents(data: List[Dict[str, Any]]) -> Tuple[List[str], List[Dict], List[str]]:
+    """
+    Prepare documents for vector store
+    
+    Args:
+        data: List of knowledge base entries
+    
+    Returns:
+        Tuple of (documents, metadatas, ids)
+    """
+    documents, metadatas, ids = [], [], []
+    
+    for item in data:
+        # Create document text with question and answer
+        document_text = f"Question: {item['Question']}\nAnswer: {item['Content']}"
+        documents.append(document_text)
+        
+        # Create metadata
+        metadata = {
+            "question": item["Question"],
+            "content": item["Content"],
+            "section": item.get("Section", ""),
+            "source": item.get("Source", ""),
+            "owner": item.get("Owner", ""),
+            "tag": item.get("Tag", ""),
+            "id": item["ID"]
+        }
+        metadatas.append(metadata)
+        
+        # Add ID
+        ids.append(item["ID"])
+    
+    return documents, metadatas, ids
+
+
+def get_knowledge_base_data() -> Tuple[List[str], List[Dict], List[str]]:
+    """
+    Load and prepare knowledge base data
+    
+    Returns:
+        Tuple of (documents, metadatas, ids)
+    """
+    df = load_knowledge_base()
+    data_list = df.to_dict('records')
+    return prepare_documents(data_list)

src/logger.py

@@ -0,0 +1,165 @@
+"""
+Logging module for XENO Bot
+Handles Google Sheets logging for responses and timing data
+"""
+import json
+import os
+from datetime import datetime
+from typing import List, Tuple, Dict, Optional
+import gspread
+from google.oauth2.service_account import Credentials
+from src.config import (
+    GOOGLE_SHEETS_CREDENTIALS_ENV,
+    SPREADSHEET_NAME,
+    RESPONSE_SHEET_INDEX,
+    TIMING_SHEET_NAME
+)
+
+
+def get_google_sheets_credentials() -> Credentials:
+    """
+    Get Google Sheets credentials from environment variable
+    
+    Returns:
+        Google Sheets credentials object
+    """
+    credentials_json = os.environ.get(GOOGLE_SHEETS_CREDENTIALS_ENV)
+    if not credentials_json:
+        raise ValueError(f"{GOOGLE_SHEETS_CREDENTIALS_ENV} environment variable not set.")
+    
+    credentials_dict = json.loads(credentials_json)
+    scope = [
+        "https://spreadsheets.google.com/feeds", 
+        "https://www.googleapis.com/auth/drive"
+    ]
+    creds = Credentials.from_service_account_info(credentials_dict, scopes=scope)
+    
+    return creds
+
+
+def initialize_sheets():
+    """
+    Initialize Google Sheets client and get sheets
+    
+    Returns:
+        Tuple of (response_sheet, timing_sheet)
+    """
+    client_gspread = gspread.authorize(get_google_sheets_credentials())
+    spreadsheet = client_gspread.open(SPREADSHEET_NAME)
+    
+    # Get response sheet
+    response_sheet = spreadsheet.get_worksheet(RESPONSE_SHEET_INDEX)
+    
+    # Get or create timing sheet
+    try:
+        timing_sheet = spreadsheet.worksheet(TIMING_SHEET_NAME)
+    except:
+        # Create timing sheet if it doesn't exist
+        timing_sheet = spreadsheet.add_worksheet(title=TIMING_SHEET_NAME, rows="1000", cols="15")
+        # Add headers
+        headers = [
+            "Timestamp", "Session_ID", "Question", "Total_Time_MS",
+            "Intent_Classification_MS", "Memory_Retrieval_MS", "RAG_Retrieval_MS", 
+            "Embedding_Generation_MS", "Similarity_Calculation_MS", "Context_Processing_MS",
+            "LLM_Generation_MS", "Memory_Update_MS", "Logging_MS", "Error_Step", "Notes"
+        ]
+        timing_sheet.append_row(headers)
+    
+    return response_sheet, timing_sheet
+
+
+# Initialize sheets
+response_sheet, timing_sheet = initialize_sheets()
+
+
+def log_response(question: str, answer: str, source_ids: str, 
+                knowledge_pairs: List[Tuple[str, str]], session_id: str, timer=None):
+    """
+    Log response to Google Sheets
+    
+    Args:
+        question: User's question
+        answer: Generated answer
+        source_ids: Source IDs used
+        knowledge_pairs: Knowledge base Q&A pairs used
+        session_id: Session identifier
+        timer: Optional timer object for tracking
+    """
+    if timer:
+        with timer.time_step("response_logging"):
+            _log_response_impl(question, answer, source_ids, knowledge_pairs, session_id)
+    else:
+        _log_response_impl(question, answer, source_ids, knowledge_pairs, session_id)
+
+
+def _log_response_impl(question: str, answer: str, source_ids: str, 
+                       knowledge_pairs: List[Tuple[str, str]], session_id: str):
+    """Internal implementation of response logging"""
+    timestamp = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
+    
+    # Extract knowledge pairs
+    knowledge_question_1 = knowledge_pairs[0][0] if len(knowledge_pairs) > 0 else "N/A"
+    knowledge_answer_1 = knowledge_pairs[0][1] if len(knowledge_pairs) > 0 else "N/A"
+    knowledge_question_2 = knowledge_pairs[1][0] if len(knowledge_pairs) > 1 else "N/A"
+    knowledge_answer_2 = knowledge_pairs[1][1] if len(knowledge_pairs) > 1 else "N/A"
+    
+    row = [
+        timestamp, session_id, question, answer, source_ids,
+        knowledge_question_1, knowledge_answer_1, 
+        knowledge_question_2, knowledge_answer_2
+    ]
+    
+    try:
+        response_sheet.append_row(row)
+        print(f"Logged response: {question} | Source IDs: {source_ids}")
+    except Exception as e:
+        print(f"Failed to log to Google Sheet: {e}")
+        # Fallback to local file
+        with open("/tmp/response_log.txt", "a") as f:
+            f.write(f"{timestamp},{question},{answer},{source_ids},{knowledge_question_1},{knowledge_answer_1},{knowledge_question_2},{knowledge_answer_2}\n")
+
+
+def log_timing_data(question: str, session_id: str, timing_summary: Dict, 
+                    error_step: Optional[str] = None, notes: Optional[str] = None):
+    """
+    Log timing data to Google Sheets
+    
+    Args:
+        question: User's question
+        session_id: Session identifier
+        timing_summary: Timing summary dictionary
+        error_step: Step where error occurred (if any)
+        notes: Additional notes
+    """
+    timestamp = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
+    step_times = timing_summary['step_times']
+    
+    # Truncate long questions
+    truncated_question = question[:100] + "..." if len(question) > 100 else question
+    
+    row = [
+        timestamp,
+        session_id,
+        truncated_question,
+        timing_summary['total_time_ms'],
+        step_times.get('intent_classification', 0),
+        step_times.get('memory_retrieval', 0),
+        step_times.get('rag_retrieval', 0),
+        step_times.get('embedding_generation', 0),
+        step_times.get('similarity_calculation', 0),
+        step_times.get('context_processing', 0),
+        step_times.get('llm_generation', 0),
+        step_times.get('memory_update', 0),
+        step_times.get('response_logging', 0),
+        error_step or "",
+        notes or ""
+    ]
+    
+    try:
+        timing_sheet.append_row(row)
+        print(f"Logged timing data: Total {timing_summary['total_time_ms']}ms")
+    except Exception as e:
+        print(f"Failed to log timing data: {e}")
+        # Fallback to local file
+        with open("/tmp/timing_log.txt", "a") as f:
+            f.write(f"{timestamp},{session_id},{question},{timing_summary}\n")

src/memory.py

@@ -0,0 +1,88 @@
+"""
+Memory module for XENO Bot
+Handles LangGraph memory operations using SQLite
+"""
+import uuid
+import sqlite3
+from datetime import datetime
+from typing import List, Dict, Any
+from langgraph.checkpoint.sqlite import SqliteSaver
+from src.config import SQLITE_DB_PATH
+
+# === LangGraph Memory Setup ===
+conn = sqlite3.connect(SQLITE_DB_PATH, check_same_thread=False)
+memory = SqliteSaver(conn=conn)
+
+
+def update_memory(config: Dict[str, Any], user_message: str, assistant_message: str, timer=None):
+    """
+    Update memory with new messages
+    
+    Args:
+        config: Configuration dictionary with thread_id
+        user_message: User's message
+        assistant_message: Assistant's response
+        timer: Optional timer object for tracking
+    """
+    if timer:
+        with timer.time_step("memory_update"):
+            _update_memory_impl(config, user_message, assistant_message)
+    else:
+        _update_memory_impl(config, user_message, assistant_message)
+
+
+def _update_memory_impl(config: Dict[str, Any], user_message: str, assistant_message: str):
+    """Internal implementation of memory update"""
+    full_checkpoint = memory.get(config) or {}
+    messages = full_checkpoint.get("channel_values", {}).get("messages", [])
+    
+    messages.append({"role": "user", "content": user_message})
+    messages.append({"role": "assistant", "content": assistant_message})
+    
+    checkpoint_to_save = {
+        "v": 1,
+        "id": str(uuid.uuid4()),
+        "ts": datetime.now().isoformat(),
+        "channel_values": {"messages": messages},
+        "channel_versions": {},
+        "versions_seen": {},   
+    }
+    
+    memory.put(config, checkpoint_to_save, {}, {})
+
+
+def retrieve_memory(config: Dict[str, Any], timer=None) -> List[Dict[str, str]]:
+    """
+    Retrieve memory messages for a session
+    
+    Args:
+        config: Configuration dictionary with thread_id
+        timer: Optional timer object for tracking
+    
+    Returns:
+        List of message dictionaries
+    """
+    if timer:
+        with timer.time_step("memory_retrieval"):
+            return _retrieve_memory_impl(config)
+    else:
+        return _retrieve_memory_impl(config)
+
+
+def _retrieve_memory_impl(config: Dict[str, Any]) -> List[Dict[str, str]]:
+    """Internal implementation of memory retrieval"""
+    full_checkpoint = memory.get(config) or {}
+    return full_checkpoint.get("channel_values", {}).get("messages", [])
+
+
+def create_session_config(session_id: str = "default") -> Dict[str, Any]:
+    """
+    Create a configuration dictionary for a session
+    
+    Args:
+        session_id: Unique session identifier
+    
+    Returns:
+        Configuration dictionary
+    """
+    return {"configurable": {"thread_id": str(session_id), "checkpoint_ns": ""}}

src/response_generator.py

@@ -0,0 +1,67 @@
+"""
+Response Generation module for XENO Bot
+Handles LLM response generation
+"""
+import google.generativeai as genai
+from typing import List, Dict
+from src.config import LLM_MODEL_NAME, SYSTEM_PROMPT
+
+
+def generate_xeno_response(context: str, question: str, chat_history: List[Dict[str, str]], timer=None) -> str:
+    """
+    Generate a response using the LLM
+    
+    Args:
+        context: Formatted context from knowledge base
+        question: User's question
+        chat_history: List of previous messages
+        timer: Optional timer object for tracking
+    
+    Returns:
+        Generated response text
+    """
+    if timer:
+        with timer.time_step("llm_generation"):
+            return _generate_response_impl(context, question, chat_history)
+    else:
+        return _generate_response_impl(context, question, chat_history)
+
+
+def _generate_response_impl(context: str, question: str, chat_history: List[Dict[str, str]]) -> str:
+    """Internal implementation of response generation"""
+    model = genai.GenerativeModel(LLM_MODEL_NAME)
+    
+    # Format chat history
+    formatted_history = "\n".join(
+        [f"{msg['role'].capitalize()}: {msg['content']}" for msg in chat_history]
+    ) if chat_history else "None"
+    
+    # Build prompt
+    prompt = f"{SYSTEM_PROMPT}\n### HISTORY ###\n{formatted_history}\n### CONTEXT ###\n{context}\n### QUESTION ###\n{question}"
+    
+    # Generate response
+    response = model.generate_content(prompt)
+    
+    return response.text.strip()
+
+
+def format_chat_history(messages: List[Dict[str, str]]) -> str:
+    """
+    Format chat history for display or logging
+    
+    Args:
+        messages: List of message dictionaries with 'role' and 'content'
+    
+    Returns:
+        Formatted string representation of chat history
+    """
+    if not messages:
+        return "No previous conversation"
+    
+    formatted = []
+    for msg in messages:
+        role = msg.get('role', 'unknown').capitalize()
+        content = msg.get('content', '')
+        formatted.append(f"{role}: {content}")
+    
+    return "\n".join(formatted)

src/utils.py

@@ -0,0 +1,66 @@
+"""
+Utilities module for XENO Bot
+Handles logging and timing functionality
+"""
+import logging
+import sys
+import time
+from contextlib import contextmanager
+from datetime import datetime
+from typing import Dict
+
+# ===== Configure Logging =====
+logging.basicConfig(
+    filename="app.log",
+    level=logging.INFO,
+    format="%(asctime)s - %(levelname)s - %(message)s"
+)
+
+def log_exception(exc_type, exc_value, exc_traceback):
+    """Log uncaught exceptions"""
+    if issubclass(exc_type, KeyboardInterrupt):
+        return
+    logging.critical("Uncaught exception", exc_info=(exc_type, exc_value, exc_traceback))
+
+sys.excepthook = log_exception
+logging.info("App started successfully.")
+
+
+# ===== Time Tracking Class =====
+class PipelineTimer:
+    """Timer for tracking pipeline execution steps"""
+    
+    def __init__(self):
+        self.reset()
+    
+    def reset(self):
+        """Reset all timing data for a new request"""
+        self.start_time = time.time()
+        self.step_times = {}
+        self.step_start = None
+        self.current_step = None
+    
+    @contextmanager
+    def time_step(self, step_name: str):
+        """Context manager to time a specific step"""
+        step_start = time.time()
+        self.current_step = step_name
+        try:
+            yield
+        finally:
+            step_end = time.time()
+            self.step_times[step_name] = round((step_end - step_start) * 1000, 2)  # Convert to milliseconds
+            self.current_step = None
+    
+    def get_total_time(self):
+        """Get total elapsed time since reset"""
+        return round((time.time() - self.start_time) * 1000, 2)
+    
+    def get_timing_summary(self) -> Dict:
+        """Get a summary of all timing data"""
+        total_time = self.get_total_time()
+        return {
+            'total_time_ms': total_time,
+            'step_times': self.step_times,
+            'timestamp': datetime.now().isoformat()
+        }

src/vector_store.py

@@ -0,0 +1,175 @@
+"""
+Vector Store module for XENO Bot
+Handles ChromaDB vector store operations
+"""
+import chromadb
+import numpy as np
+import torch
+from langchain_chroma import Chroma
+from sentence_transformers import util
+from typing import List, Tuple, Any
+import google.generativeai as genai
+from src.config import (
+    COLLECTION_NAME, 
+    CHROMA_DB_PATH, 
+    RAG_TOP_K, 
+    RAG_MAX_RESULTS,
+    EMBEDDING_MODEL
+)
+from src.knowledge_base import get_knowledge_base_data
+
+
+def initialize_vector_store() -> Tuple[chromadb.Collection, Chroma, Any]:
+    """
+    Initialize ChromaDB vector store
+    
+    Returns:
+        Tuple of (collection, vector_store, retriever)
+    """
+    # Get knowledge base data
+    documents, metadatas, ids = get_knowledge_base_data()
+    
+    # Initialize ChromaDB client
+    try:
+        client = chromadb.PersistentClient(path=CHROMA_DB_PATH)
+        
+        # Try to get existing collection
+        try:
+            collection = client.get_collection(name=COLLECTION_NAME)
+            print(f"Loaded existing ChromaDB collection: {COLLECTION_NAME}")
+        except:
+            # Create new collection if it doesn't exist
+            print(f"Creating new ChromaDB collection: {COLLECTION_NAME}")
+            collection = client.create_collection(name=COLLECTION_NAME)
+            collection.add(documents=documents, metadatas=metadatas, ids=ids)
+        
+        # Create vector store and retriever
+        vector_store = Chroma(client=client, collection_name=COLLECTION_NAME)
+        retriever = vector_store.as_retriever(
+            search_type="similarity", 
+            search_kwargs={"k": RAG_TOP_K}
+        )
+        
+        return collection, vector_store, retriever
+        
+    except Exception as e:
+        print(f"Failed to initialize ChromaDB: {e}")
+        raise
+
+
+def generate_embeddings(query: str, documents: List[Any], timer=None) -> Tuple[List[float], List[List[float]]]:
+    """
+    Generate embeddings for query and documents
+    
+    Args:
+        query: User query
+        documents: List of retrieved documents
+        timer: Optional timer object for tracking
+    
+    Returns:
+        Tuple of (query_embedding, doc_embeddings)
+    """
+    if timer:
+        with timer.time_step("embedding_generation"):
+            return _generate_embeddings_impl(query, documents)
+    else:
+        return _generate_embeddings_impl(query, documents)
+
+
+def _generate_embeddings_impl(query: str, documents: List[Any]) -> Tuple[List[float], List[List[float]]]:
+    """Internal implementation of embedding generation"""
+    # Generate query embedding
+    query_embedding = genai.embed_content(
+        model=EMBEDDING_MODEL, 
+        content=query, 
+        task_type="retrieval_query"
+    )['embedding']
+    
+    # Generate document embeddings
+    doc_embeddings = [
+        genai.embed_content(
+            model=EMBEDDING_MODEL, 
+            content=doc.page_content, 
+            task_type="retrieval_document"
+        )['embedding'] 
+        for doc in documents
+    ]
+    
+    return query_embedding, doc_embeddings
+
+
+def calculate_similarity(query_embedding: List[float], doc_embeddings: List[List[float]], timer=None) -> List[float]:
+    """
+    Calculate cosine similarity between query and documents
+    
+    Args:
+        query_embedding: Query embedding vector
+        doc_embeddings: List of document embedding vectors
+        timer: Optional timer object for tracking
+    
+    Returns:
+        List of cosine similarity scores
+    """
+    if timer:
+        with timer.time_step("similarity_calculation"):
+            return _calculate_similarity_impl(query_embedding, doc_embeddings)
+    else:
+        return _calculate_similarity_impl(query_embedding, doc_embeddings)
+
+
+def _calculate_similarity_impl(query_embedding: List[float], doc_embeddings: List[List[float]]) -> List[float]:
+    """Internal implementation of similarity calculation"""
+    cosine_scores = util.cos_sim(
+        torch.tensor(query_embedding).float(), 
+        torch.tensor(doc_embeddings).float()
+    )[0].tolist()
+    
+    return cosine_scores
+
+
+def process_context(results: List[Any], cosine_scores: List[float], 
+                    max_results: int = RAG_MAX_RESULTS, timer=None) -> Tuple[str, List[str], List[Tuple[str, str]]]:
+    """
+    Process retrieved context and format for LLM
+    
+    Args:
+        results: List of retrieved documents
+        cosine_scores: List of similarity scores
+        max_results: Maximum number of results to include
+        timer: Optional timer object for tracking
+    
+    Returns:
+        Tuple of (formatted_context, source_ids, knowledge_pairs)
+    """
+    if timer:
+        with timer.time_step("context_processing"):
+            return _process_context_impl(results, cosine_scores, max_results)
+    else:
+        return _process_context_impl(results, cosine_scores, max_results)
+
+
+def _process_context_impl(results: List[Any], cosine_scores: List[float], 
+                          max_results: int) -> Tuple[str, List[str], List[Tuple[str, str]]]:
+    """Internal implementation of context processing"""
+    sorted_indices = np.argsort(cosine_scores)[::-1][:max_results]
+    
+    formatted_context = ""
+    source_ids = []
+    knowledge_pairs = []
+    
+    for i, idx in enumerate(sorted_indices, 1):
+        result = results[idx]
+        score = cosine_scores[idx]
+        
+        question = result.metadata.get('question', 'N/A')
+        answer = result.metadata.get('content', 'N/A')
+        
+        formatted_context += f"Knowledge Entry {i}:\n"
+        formatted_context += f"Q: {question}\n"
+        formatted_context += f"A: {answer}\n"
+        formatted_context += "-" * 40 + "\n"
+        
+        source_ids.append(result.metadata.get('id', 'N/A'))
+        knowledge_pairs.append((question, answer))
+    
+    return formatted_context, source_ids, knowledge_pairs

tests/README.md

@@ -0,0 +1,119 @@
+# XENO Bot - Unit Tests
+
+This directory contains comprehensive unit tests for the XENO Bot application.
+
+## Test Coverage
+
+The test suite covers the following modules:
+
+1. **test_utils.py** - Tests for PipelineTimer and logging utilities
+2. **test_intent_classifier.py** - Tests for intent classification (greetings, thanks, goodbye, queries)
+3. **test_knowledge_base.py** - Tests for knowledge base loading and document preparation
+4. **test_memory.py** - Tests for LangGraph memory operations (SQLite-based)
+5. **test_response_generator.py** - Tests for LLM response generation
+6. **test_vector_store.py** - Tests for ChromaDB vector store operations
+7. **test_logger.py** - Tests for Google Sheets logging functionality
+
+## Running Tests
+
+### Install Test Dependencies
+
+```bash
+pip install -r requirements.txt
+```
+
+### Run All Tests
+
+```bash
+pytest
+```
+
+### Run Specific Test File
+
+```bash
+pytest tests/test_utils.py
+```
+
+### Run with Coverage Report
+
+```bash
+pytest --cov=src --cov-report=html
+```
+
+This will generate an HTML coverage report in the `htmlcov/` directory.
+
+### Run with Verbose Output
+
+```bash
+pytest -v
+```
+
+### Run Specific Test
+
+```bash
+pytest tests/test_intent_classifier.py::TestIntentClassifier::test_classify_greeting
+```
+
+## Test Structure
+
+Each test file follows a consistent structure:
+
+- **setUp()**: Initialize test fixtures and mock data
+- **tearDown()**: Clean up after tests (if needed)
+- **test_***: Individual test methods
+
+## Mocking
+
+Tests use Python's `unittest.mock` to:
+- Mock external API calls (Google Generative AI)
+- Mock database connections (ChromaDB, SQLite)
+- Mock Google Sheets operations
+- Isolate units under test
+
+## Best Practices
+
+- Tests are isolated and don't depend on external services
+- Each test focuses on a single behavior
+- Mock objects are used to simulate dependencies
+- Tests include both positive and negative scenarios
+- Edge cases are covered (empty inputs, errors, etc.)
+
+## Continuous Integration
+
+To integrate with CI/CD pipelines, add to your workflow:
+
+```yaml
+- name: Run tests
+  run: |
+    pip install -r requirements.txt
+    pytest --cov=src --cov-report=xml
+```
+
+## Test Configuration
+
+Test configuration is defined in:
+- `pytest.ini` - Pytest configuration
+- `setup.cfg` - Additional pytest and coverage settings
+
+## Writing New Tests
+
+When adding new functionality:
+
+1. Create a test file in `tests/` following the naming convention `test_<module>.py`
+2. Create test class inheriting from `unittest.TestCase`
+3. Add test methods starting with `test_`
+4. Use mocks for external dependencies
+5. Run tests to ensure they pass
+
+Example:
+
+```python
+import unittest
+from unittest.mock import patch
+from src.my_module import my_function
+
+class TestMyModule(unittest.TestCase):
+    def test_my_function(self):
+        result = my_function("test")
+        self.assertEqual(result, expected_value)
+```

tests/__init__.py

@@ -0,0 +1,3 @@
+"""
+Unit tests for XENO Bot
+"""

tests/conftest.py

@@ -0,0 +1,117 @@
+"""
+Pytest configuration file
+Sets up test environment and fixtures
+"""
+import os
+import sys
+import pytest
+from unittest.mock import Mock, MagicMock, patch, PropertyMock
+
+# Add src to path
+sys.path.insert(0, os.path.join(os.path.dirname(__file__), '..'))
+
+# Set mock environment variables before importing any modules
+os.environ.setdefault('GEMINI_API_KEY', 'test-api-key-12345')
+
+# Mock Google Sheets credentials
+mock_credentials = {
+    "type": "service_account",
+    "project_id": "test-project",
+    "private_key_id": "test-key-id",
+    "private_key": "-----BEGIN PRIVATE KEY-----\ntest\n-----END PRIVATE KEY-----\n",
+    "client_email": "test@test.iam.gserviceaccount.com",
+    "client_id": "12345",
+    "auth_uri": "https://accounts.google.com/o/oauth2/auth",
+    "token_uri": "https://oauth2.googleapis.com/token",
+    "auth_provider_x509_cert_url": "https://www.googleapis.com/oauth2/v1/certs",
+    "client_x509_cert_url": "https://www.googleapis.com/robot/v1/metadata/x509/test"
+}
+import json
+os.environ.setdefault('GOOGLE_SHEETS_CREDENTIALS', json.dumps(mock_credentials))
+
+# Mock google.oauth2 and gspread modules before src.logger imports them
+mock_credentials_class = MagicMock()
+mock_creds_instance = MagicMock()
+mock_credentials_class.from_service_account_info = Mock(return_value=mock_creds_instance)
+
+mock_oauth2 = MagicMock()
+mock_oauth2.service_account.Credentials = mock_credentials_class
+sys.modules['google.oauth2'] = mock_oauth2
+sys.modules['google.oauth2.service_account'] = mock_oauth2.service_account
+
+mock_gspread = MagicMock()
+mock_spreadsheet = MagicMock()
+mock_worksheet = MagicMock()
+mock_worksheet.append_row = Mock()
+mock_spreadsheet.get_worksheet = Mock(return_value=mock_worksheet)
+mock_spreadsheet.worksheet = Mock(return_value=mock_worksheet)
+mock_spreadsheet.add_worksheet = Mock(return_value=mock_worksheet)
+mock_client = MagicMock()
+mock_client.open = Mock(return_value=mock_spreadsheet)
+mock_gspread.authorize = Mock(return_value=mock_client)
+sys.modules['gspread'] = mock_gspread
+
+
+@pytest.fixture(autouse=True)
+def mock_google_sheets():
+    """Mock Google Sheets to avoid actual connections during testing"""
+    with patch('src.logger.response_sheet') as mock_response, \
+         patch('src.logger.timing_sheet') as mock_timing:
+        mock_response.append_row = Mock()
+        mock_timing.append_row = Mock()
+        yield mock_response, mock_timing
+
+
+@pytest.fixture
+def mock_genai():
+    """Mock Google Generative AI"""
+    with patch('google.generativeai.configure') as mock_config, \
+         patch('google.generativeai.GenerativeModel') as mock_model, \
+         patch('google.generativeai.embed_content') as mock_embed:
+        yield {
+            'configure': mock_config,
+            'model': mock_model,
+            'embed': mock_embed
+        }
+
+
+@pytest.fixture
+def mock_chromadb():
+    """Mock ChromaDB client"""
+    with patch('chromadb.PersistentClient') as mock_client:
+        mock_collection = Mock()
+        mock_client.return_value.get_collection.return_value = mock_collection
+        yield mock_client
+
+
+@pytest.fixture
+def mock_sqlite():
+    """Mock SQLite connections for memory"""
+    with patch('sqlite3.connect') as mock_connect:
+        mock_conn = Mock()
+        mock_connect.return_value = mock_conn
+        yield mock_conn
+
+
+@pytest.fixture
+def sample_documents():
+    """Provide sample documents for testing"""
+    doc1 = Mock()
+    doc1.page_content = "Question: How do I create an account?\nAnswer: Visit our website."
+    doc1.metadata = {
+        'id': 'KB001',
+        'question': 'How do I create an account?',
+        'content': 'Visit our website.',
+        'section': 'Account Management'
+    }
+    
+    doc2 = Mock()
+    doc2.page_content = "Question: What are the fees?\nAnswer: 1% per transaction."
+    doc2.metadata = {
+        'id': 'KB002',
+        'question': 'What are the fees?',
+        'content': '1% per transaction.',
+        'section': 'Fees'
+    }
+    
+    return [doc1, doc2]

tests/test_intent_classifier.py

@@ -0,0 +1,169 @@
+"""
+Unit tests for intent_classifier module
+Tests the IntentClassifier class
+"""
+import unittest
+from unittest.mock import Mock
+from src.intent_classifier import IntentClassifier
+
+
+class TestIntentClassifier(unittest.TestCase):
+    """Test cases for IntentClassifier class"""
+    
+    def setUp(self):
+        """Set up test fixtures"""
+        self.classifier = IntentClassifier()
+    
+    def test_initialization(self):
+        """Test classifier initialization"""
+        self.assertIsNotNone(self.classifier.intent_patterns)
+        self.assertIn('greeting', self.classifier.intent_patterns)
+        self.assertIn('thanks', self.classifier.intent_patterns)
+        self.assertIn('goodbye', self.classifier.intent_patterns)
+    
+    def test_classify_greeting(self):
+        """Test classification of greeting messages"""
+        test_cases = [
+            "hi",
+            "hello",
+            "Hey there",
+            "good morning",
+            "Good afternoon!",
+            "how are you"
+        ]
+        
+        for message in test_cases:
+            intent, response = self.classifier.classify_intent(message)
+            self.assertEqual(intent, 'greeting', f"Failed for message: {message}")
+            self.assertIsInstance(response, str)
+            self.assertGreater(len(response), 0)
+    
+    def test_classify_thanks(self):
+        """Test classification of thank you messages"""
+        test_cases = [
+            "thank you",
+            "thanks",
+            "thank u",
+            "thx",
+            "I appreciate it",
+            "thanks a lot",
+            "thank you so much"
+        ]
+        
+        for message in test_cases:
+            intent, response = self.classifier.classify_intent(message)
+            self.assertEqual(intent, 'thanks', f"Failed for message: {message}")
+            self.assertIsInstance(response, str)
+            self.assertGreater(len(response), 0)
+    
+    def test_classify_goodbye(self):
+        """Test classification of goodbye messages"""
+        test_cases = [
+            "bye",
+            "goodbye",
+            "see you",
+            "farewell",
+            "take care",
+            "have a good day",
+            "talk to you later"
+        ]
+        
+        for message in test_cases:
+            intent, response = self.classifier.classify_intent(message)
+            self.assertEqual(intent, 'goodbye', f"Failed for message: {message}")
+            self.assertIsInstance(response, str)
+            self.assertGreater(len(response), 0)
+    
+    def test_classify_query(self):
+        """Test classification of query messages"""
+        test_cases = [
+            "How do I open an account?",
+            "What are the transaction fees?",
+            "Can you help me with my balance?",
+            "Tell me about XENO services"
+        ]
+        
+        for message in test_cases:
+            intent, response = self.classifier.classify_intent(message)
+            self.assertEqual(intent, 'query', f"Failed for message: {message}")
+            self.assertEqual(response, '')
+    
+    def test_case_insensitivity(self):
+        """Test that classification is case insensitive"""
+        messages = [
+            ("HI", 'greeting'),
+            ("THANK YOU", 'thanks'),
+            ("BYE", 'goodbye'),
+            ("Hi There", 'greeting')
+        ]
+        
+        for message, expected_intent in messages:
+            intent, _ = self.classifier.classify_intent(message)
+            self.assertEqual(intent, expected_intent)
+    
+    def test_with_timer(self):
+        """Test classification with timer object"""
+        mock_timer = Mock()
+        mock_timer.time_step = Mock()
+        mock_timer.time_step.return_value.__enter__ = Mock()
+        mock_timer.time_step.return_value.__exit__ = Mock()
+        
+        intent, response = self.classifier.classify_intent("hello", timer=mock_timer)
+        
+        self.assertEqual(intent, 'greeting')
+        mock_timer.time_step.assert_called_once_with("intent_classification")
+    
+    def test_is_simple_intent(self):
+        """Test is_simple_intent method"""
+        self.assertTrue(self.classifier.is_simple_intent('greeting'))
+        self.assertTrue(self.classifier.is_simple_intent('thanks'))
+        self.assertFalse(self.classifier.is_simple_intent('goodbye'))
+        self.assertFalse(self.classifier.is_simple_intent('query'))
+    
+    def test_add_intent(self):
+        """Test adding a new intent"""
+        patterns = [r'\b(test|testing)\b']
+        responses = ["This is a test response"]
+        
+        self.classifier.add_intent('test_intent', patterns, responses)
+        
+        # Verify intent was added
+        self.assertIn('test_intent', self.classifier.intent_patterns)
+        self.assertEqual(self.classifier.intent_patterns['test_intent']['patterns'], patterns)
+        self.assertEqual(self.classifier.intent_patterns['test_intent']['responses'], responses)
+        
+        # Test classification with new intent
+        intent, response = self.classifier.classify_intent("testing")
+        self.assertEqual(intent, 'test_intent')
+        self.assertEqual(response, "This is a test response")
+    
+    def test_response_variety(self):
+        """Test that responses vary (random selection)"""
+        # Multiple calls might return different responses
+        responses = set()
+        for _ in range(20):
+            _, response = self.classifier.classify_intent("hello")
+            responses.add(response)
+        
+        # Should have at least 1 response (could be more if random varies)
+        self.assertGreater(len(responses), 0)
+    
+    def test_empty_message(self):
+        """Test classification of empty or whitespace messages"""
+        test_cases = ["", "   ", "\n", "\t"]
+        
+        for message in test_cases:
+            intent, response = self.classifier.classify_intent(message)
+            self.assertEqual(intent, 'query')
+            self.assertEqual(response, '')
+    
+    def test_mixed_intent_message(self):
+        """Test messages that might match multiple patterns"""
+        # "hi thank you" should match greeting (first match wins)
+        intent, response = self.classifier.classify_intent("hi thank you")
+        # Should match the first pattern it encounters
+        self.assertIn(intent, ['greeting', 'thanks'])
+
+
+if __name__ == '__main__':
+    unittest.main()

tests/test_knowledge_base.py

@@ -0,0 +1,183 @@
+"""
+Unit tests for knowledge_base module
+Tests knowledge base loading and preparation
+"""
+import unittest
+import pandas as pd
+import json
+import tempfile
+import os
+from unittest.mock import patch, Mock
+from src.knowledge_base import (
+    load_knowledge_base,
+    prepare_documents,
+    get_knowledge_base_data
+)
+
+
+class TestKnowledgeBase(unittest.TestCase):
+    """Test cases for knowledge_base module"""
+    
+    def setUp(self):
+        """Set up test fixtures"""
+        # Create sample knowledge base data
+        self.sample_data = [
+            {
+                "ID": "KB001",
+                "Question": "How do I create an account?",
+                "Content": "You can create an account by visiting our website.",
+                "Section": "Account Management",
+                "Source": "Website",
+                "Owner": "Support Team",
+                "Tag": "account"
+            },
+            {
+                "ID": "KB002",
+                "Question": "What are the fees?",
+                "Content": "Our transaction fees are 1% per transaction.",
+                "Section": "Fees",
+                "Source": "Documentation",
+                "Owner": "Finance Team",
+                "Tag": "fees"
+            }
+        ]
+        
+        # Create temporary JSON file
+        self.temp_file = tempfile.NamedTemporaryFile(
+            mode='w', 
+            delete=False, 
+            suffix='.json'
+        )
+        json.dump(self.sample_data, self.temp_file)
+        self.temp_file.close()
+    
+    def tearDown(self):
+        """Clean up test fixtures"""
+        if os.path.exists(self.temp_file.name):
+            os.unlink(self.temp_file.name)
+    
+    def test_load_knowledge_base(self):
+        """Test loading knowledge base from JSON file"""
+        df = load_knowledge_base(self.temp_file.name)
+        
+        # Check DataFrame structure
+        self.assertIsInstance(df, pd.DataFrame)
+        self.assertEqual(len(df), 2)
+        self.assertIn('ID', df.columns)
+        self.assertIn('Question', df.columns)
+        self.assertIn('Content', df.columns)
+    
+    def test_load_knowledge_base_drops_null_content(self):
+        """Test that rows with null Content are dropped"""
+        data_with_null = self.sample_data + [
+            {
+                "ID": "KB003",
+                "Question": "Test question?",
+                "Content": None,
+                "Section": "Test"
+            }
+        ]
+        
+        temp_file_null = tempfile.NamedTemporaryFile(
+            mode='w', 
+            delete=False, 
+            suffix='.json'
+        )
+        json.dump(data_with_null, temp_file_null)
+        temp_file_null.close()
+        
+        try:
+            df = load_knowledge_base(temp_file_null.name)
+            # Should only have 2 rows (null Content row dropped)
+            self.assertEqual(len(df), 2)
+        finally:
+            os.unlink(temp_file_null.name)
+    
+    def test_prepare_documents(self):
+        """Test preparing documents for vector store"""
+        documents, metadatas, ids = prepare_documents(self.sample_data)
+        
+        # Check lengths match
+        self.assertEqual(len(documents), 2)
+        self.assertEqual(len(metadatas), 2)
+        self.assertEqual(len(ids), 2)
+        
+        # Check document format
+        self.assertIn("Question:", documents[0])
+        self.assertIn("Answer:", documents[0])
+        self.assertIn("How do I create an account?", documents[0])
+        
+        # Check metadata structure
+        self.assertEqual(metadatas[0]['id'], 'KB001')
+        self.assertEqual(metadatas[0]['question'], 'How do I create an account?')
+        self.assertEqual(metadatas[0]['section'], 'Account Management')
+        
+        # Check IDs
+        self.assertEqual(ids[0], 'KB001')
+        self.assertEqual(ids[1], 'KB002')
+    
+    def test_prepare_documents_with_missing_fields(self):
+        """Test preparing documents with missing optional fields"""
+        data_minimal = [
+            {
+                "ID": "KB001",
+                "Question": "Test question?",
+                "Content": "Test answer."
+            }
+        ]
+        
+        documents, metadatas, ids = prepare_documents(data_minimal)
+        
+        # Should still work with defaults
+        self.assertEqual(len(documents), 1)
+        self.assertEqual(metadatas[0]['section'], '')
+        self.assertEqual(metadatas[0]['source'], '')
+        self.assertEqual(metadatas[0]['owner'], '')
+        self.assertEqual(metadatas[0]['tag'], '')
+    
+    @patch('src.knowledge_base.load_knowledge_base')
+    def test_get_knowledge_base_data(self, mock_load):
+        """Test get_knowledge_base_data function"""
+        # Mock the load_knowledge_base function
+        mock_df = pd.DataFrame(self.sample_data)
+        mock_load.return_value = mock_df
+        
+        documents, metadatas, ids = get_knowledge_base_data()
+        
+        # Verify load was called
+        mock_load.assert_called_once()
+        
+        # Verify output
+        self.assertEqual(len(documents), 2)
+        self.assertEqual(len(metadatas), 2)
+        self.assertEqual(len(ids), 2)
+    
+    def test_document_text_format(self):
+        """Test that document text is properly formatted"""
+        documents, _, _ = prepare_documents(self.sample_data)
+        
+        # Check first document format
+        expected_format = "Question: How do I create an account?\nAnswer: You can create an account by visiting our website."
+        self.assertEqual(documents[0], expected_format)
+    
+    def test_empty_knowledge_base(self):
+        """Test handling of empty knowledge base"""
+        empty_data = []
+        documents, metadatas, ids = prepare_documents(empty_data)
+        
+        self.assertEqual(len(documents), 0)
+        self.assertEqual(len(metadatas), 0)
+        self.assertEqual(len(ids), 0)
+    
+    def test_metadata_completeness(self):
+        """Test that all metadata fields are present"""
+        _, metadatas, _ = prepare_documents(self.sample_data)
+        
+        required_fields = ['question', 'content', 'section', 'source', 'owner', 'tag', 'id']
+        for metadata in metadatas:
+            for field in required_fields:
+                self.assertIn(field, metadata)
+
+
+if __name__ == '__main__':
+    unittest.main()

tests/test_logger.py

@@ -0,0 +1,272 @@
+"""
+Unit tests for logger module
+Tests Google Sheets logging functionality
+"""
+import unittest
+from datetime import datetime
+from unittest.mock import patch, Mock, MagicMock
+from src.logger import (
+    log_response,
+    log_timing_data,
+    _log_response_impl
+)
+
+
+class TestLogger(unittest.TestCase):
+    """Test cases for logger module"""
+    
+    def setUp(self):
+        """Set up test fixtures"""
+        self.question = "How do I create an account?"
+        self.answer = "You can create an account by visiting our website."
+        self.source_ids = "KB001, KB002"
+        self.knowledge_pairs = [
+            ("Question 1?", "Answer 1."),
+            ("Question 2?", "Answer 2.")
+        ]
+        self.session_id = "test_session_123"
+    
+    @patch('src.logger.response_sheet')
+    def test_log_response_impl(self, mock_sheet):
+        """Test internal response logging implementation"""
+        _log_response_impl(
+            self.question,
+            self.answer,
+            self.source_ids,
+            self.knowledge_pairs,
+            self.session_id
+        )
+        
+        # Verify append_row was called
+        mock_sheet.append_row.assert_called_once()
+        
+        # Check the row data
+        call_args = mock_sheet.append_row.call_args
+        row = call_args[0][0]
+        
+        # Verify row structure
+        self.assertEqual(len(row), 9)  # timestamp, session_id, question, answer, source_ids, 4 knowledge fields
+        self.assertEqual(row[1], self.session_id)
+        self.assertEqual(row[2], self.question)
+        self.assertEqual(row[3], self.answer)
+        self.assertEqual(row[4], self.source_ids)
+        self.assertEqual(row[5], "Question 1?")
+        self.assertEqual(row[6], "Answer 1.")
+        self.assertEqual(row[7], "Question 2?")
+        self.assertEqual(row[8], "Answer 2.")
+    
+    @patch('src.logger.response_sheet')
+    def test_log_response_with_timer(self, mock_sheet):
+        """Test log_response with timer"""
+        mock_timer = Mock()
+        mock_timer.time_step = MagicMock()
+        mock_timer.time_step.return_value.__enter__ = Mock()
+        mock_timer.time_step.return_value.__exit__ = Mock()
+        
+        log_response(
+            self.question,
+            self.answer,
+            self.source_ids,
+            self.knowledge_pairs,
+            self.session_id,
+            timer=mock_timer
+        )
+        
+        # Verify timer was used
+        mock_timer.time_step.assert_called_once_with("response_logging")
+    
+    @patch('src.logger.response_sheet')
+    def test_log_response_empty_knowledge_pairs(self, mock_sheet):
+        """Test logging with empty knowledge pairs"""
+        _log_response_impl(
+            self.question,
+            self.answer,
+            self.source_ids,
+            [],
+            self.session_id
+        )
+        
+        # Should still work
+        mock_sheet.append_row.assert_called_once()
+        
+        # Check that N/A is used for missing pairs
+        row = mock_sheet.append_row.call_args[0][0]
+        self.assertEqual(row[5], "N/A")
+        self.assertEqual(row[6], "N/A")
+    
+    @patch('src.logger.response_sheet')
+    def test_log_response_single_knowledge_pair(self, mock_sheet):
+        """Test logging with single knowledge pair"""
+        single_pair = [("Single question?", "Single answer.")]
+        
+        _log_response_impl(
+            self.question,
+            self.answer,
+            self.source_ids,
+            single_pair,
+            self.session_id
+        )
+        
+        row = mock_sheet.append_row.call_args[0][0]
+        
+        # First pair should be present
+        self.assertEqual(row[5], "Single question?")
+        self.assertEqual(row[6], "Single answer.")
+        
+        # Second pair should be N/A
+        self.assertEqual(row[7], "N/A")
+        self.assertEqual(row[8], "N/A")
+    
+    @patch('src.logger.response_sheet')
+    @patch('builtins.open', create=True)
+    def test_log_response_fallback_on_error(self, mock_open, mock_sheet):
+        """Test fallback to file logging on error"""
+        # Make append_row raise an exception
+        mock_sheet.append_row.side_effect = Exception("Connection error")
+        
+        # Mock file operations
+        mock_file = MagicMock()
+        mock_open.return_value.__enter__.return_value = mock_file
+        
+        # Should not raise exception
+        _log_response_impl(
+            self.question,
+            self.answer,
+            self.source_ids,
+            self.knowledge_pairs,
+            self.session_id
+        )
+        
+        # Verify fallback file was opened
+        mock_open.assert_called_once_with("/tmp/response_log.txt", "a")
+        mock_file.write.assert_called_once()
+    
+    @patch('src.logger.timing_sheet')
+    def test_log_timing_data(self, mock_sheet):
+        """Test timing data logging"""
+        timing_summary = {
+            'total_time_ms': 1500,
+            'step_times': {
+                'intent_classification': 50,
+                'memory_retrieval': 100,
+                'rag_retrieval': 200,
+                'embedding_generation': 300,
+                'similarity_calculation': 150,
+                'context_processing': 100,
+                'llm_generation': 500,
+                'memory_update': 50,
+                'response_logging': 50
+            }
+        }
+        
+        log_timing_data(
+            self.question,
+            self.session_id,
+            timing_summary,
+            error_step=None,
+            notes="Test note"
+        )
+        
+        # Verify append_row was called
+        mock_sheet.append_row.assert_called_once()
+        
+        # Check row structure
+        row = mock_sheet.append_row.call_args[0][0]
+        
+        # Should have 15 fields
+        self.assertEqual(len(row), 15)
+        self.assertEqual(row[1], self.session_id)
+        self.assertEqual(row[3], 1500)  # total_time_ms
+        self.assertEqual(row[4], 50)    # intent_classification
+        self.assertEqual(row[5], 100)   # memory_retrieval
+        self.assertEqual(row[14], "Test note")  # notes
+    
+    @patch('src.logger.timing_sheet')
+    def test_log_timing_data_with_error(self, mock_sheet):
+        """Test timing data logging with error"""
+        timing_summary = {
+            'total_time_ms': 500,
+            'step_times': {
+                'intent_classification': 50
+            }
+        }
+        
+        log_timing_data(
+            self.question,
+            self.session_id,
+            timing_summary,
+            error_step="rag_retrieval",
+            notes="Error occurred"
+        )
+        
+        row = mock_sheet.append_row.call_args[0][0]
+        
+        # Check error_step is logged
+        self.assertEqual(row[13], "rag_retrieval")
+        self.assertEqual(row[14], "Error occurred")
+    
+    @patch('src.logger.timing_sheet')
+    def test_log_timing_data_missing_steps(self, mock_sheet):
+        """Test timing data with missing step times"""
+        timing_summary = {
+            'total_time_ms': 100,
+            'step_times': {
+                'intent_classification': 100
+                # Other steps missing
+            }
+        }
+        
+        log_timing_data(
+            self.question,
+            self.session_id,
+            timing_summary
+        )
+        
+        row = mock_sheet.append_row.call_args[0][0]
+        
+        # Missing steps should default to 0
+        self.assertEqual(row[5], 0)  # memory_retrieval
+        self.assertEqual(row[6], 0)  # rag_retrieval
+    
+    @patch('src.logger.timing_sheet')
+    def test_log_timing_data_long_question(self, mock_sheet):
+        """Test timing data logging with long question (truncation)"""
+        long_question = "A" * 150  # 150 characters
+        
+        timing_summary = {
+            'total_time_ms': 100,
+            'step_times': {}
+        }
+        
+        log_timing_data(
+            long_question,
+            self.session_id,
+            timing_summary
+        )
+        
+        row = mock_sheet.append_row.call_args[0][0]
+        
+        # Question should be truncated to 103 chars (100 + "...")
+        self.assertEqual(len(row[2]), 103)
+        self.assertTrue(row[2].endswith("..."))
+    
+    @patch('src.logger.timing_sheet')
+    @patch('builtins.open', create=True)
+    def test_log_timing_data_fallback_on_error(self, mock_open, mock_sheet):
+        """Test fallback to file logging for timing data on error"""
+        mock_sheet.append_row.side_effect = Exception("Connection error")
+        
+        mock_file = MagicMock()
+        mock_open.return_value.__enter__.return_value = mock_file
+        
+        timing_summary = {'total_time_ms': 100, 'step_times': {}}
+        
+        log_timing_data(self.question, self.session_id, timing_summary)
+        
+        # Verify fallback file was opened
+        mock_open.assert_called_once_with("/tmp/timing_log.txt", "a")
+        mock_file.write.assert_called_once()
+
+
+if __name__ == '__main__':
+    unittest.main()

tests/test_memory.py

@@ -0,0 +1,195 @@
+"""
+Unit tests for memory module
+Tests LangGraph memory operations
+"""
+import unittest
+import os
+import sqlite3
+import tempfile
+from unittest.mock import patch, Mock, MagicMock
+from src.memory import (
+    update_memory,
+    retrieve_memory,
+    create_session_config,
+    _update_memory_impl,
+    _retrieve_memory_impl
+)
+
+
+class TestMemory(unittest.TestCase):
+    """Test cases for memory module"""
+    
+    def setUp(self):
+        """Set up test fixtures"""
+        self.test_config = {
+            "configurable": {
+                "thread_id": "test_session_123",
+                "checkpoint_ns": ""
+            }
+        }
+    
+    def test_create_session_config(self):
+        """Test creating session config"""
+        session_id = "test_session_456"
+        config = create_session_config(session_id)
+        
+        # Check structure
+        self.assertIn("configurable", config)
+        self.assertEqual(config["configurable"]["thread_id"], session_id)
+        self.assertEqual(config["configurable"]["checkpoint_ns"], "")
+    
+    def test_create_session_config_default(self):
+        """Test creating session config with default ID"""
+        config = create_session_config()
+        
+        # Check structure
+        self.assertIn("configurable", config)
+        self.assertEqual(config["configurable"]["thread_id"], "default")
+    
+    @patch('src.memory.memory')
+    def test_update_memory_impl(self, mock_memory):
+        """Test internal memory update implementation"""
+        # Mock memory.get to return existing checkpoint
+        mock_checkpoint = {
+            "channel_values": {
+                "messages": [
+                    {"role": "user", "content": "Previous question"},
+                    {"role": "assistant", "content": "Previous answer"}
+                ]
+            }
+        }
+        mock_memory.get.return_value = mock_checkpoint
+        
+        user_message = "New question"
+        assistant_message = "New answer"
+        
+        _update_memory_impl(self.test_config, user_message, assistant_message)
+        
+        # Verify memory.get was called
+        mock_memory.get.assert_called_once_with(self.test_config)
+        
+        # Verify memory.put was called
+        mock_memory.put.assert_called_once()
+        
+        # Check the checkpoint that was saved
+        call_args = mock_memory.put.call_args
+        saved_checkpoint = call_args[0][1]
+        
+        # Verify messages were appended
+        messages = saved_checkpoint["channel_values"]["messages"]
+        self.assertEqual(len(messages), 4)  # 2 existing + 2 new
+        self.assertEqual(messages[-2]["role"], "user")
+        self.assertEqual(messages[-2]["content"], user_message)
+        self.assertEqual(messages[-1]["role"], "assistant")
+        self.assertEqual(messages[-1]["content"], assistant_message)
+    
+    @patch('src.memory.memory')
+    def test_update_memory_empty_checkpoint(self, mock_memory):
+        """Test updating memory with empty checkpoint"""
+        # Mock memory.get to return None
+        mock_memory.get.return_value = None
+        
+        user_message = "First question"
+        assistant_message = "First answer"
+        
+        _update_memory_impl(self.test_config, user_message, assistant_message)
+        
+        # Verify memory.put was called
+        mock_memory.put.assert_called_once()
+        
+        # Check the checkpoint
+        call_args = mock_memory.put.call_args
+        saved_checkpoint = call_args[0][1]
+        messages = saved_checkpoint["channel_values"]["messages"]
+        
+        # Should have 2 messages
+        self.assertEqual(len(messages), 2)
+        self.assertEqual(messages[0]["role"], "user")
+        self.assertEqual(messages[1]["role"], "assistant")
+    
+    @patch('src.memory.memory')
+    def test_update_memory_with_timer(self, mock_memory):
+        """Test update_memory with timer"""
+        mock_memory.get.return_value = {}
+        mock_timer = Mock()
+        mock_timer.time_step = MagicMock()
+        mock_timer.time_step.return_value.__enter__ = Mock()
+        mock_timer.time_step.return_value.__exit__ = Mock()
+        
+        update_memory(self.test_config, "Test", "Answer", timer=mock_timer)
+        
+        # Verify timer was used
+        mock_timer.time_step.assert_called_once_with("memory_update")
+    
+    @patch('src.memory.memory')
+    def test_retrieve_memory_impl(self, mock_memory):
+        """Test internal memory retrieval implementation"""
+        # Mock memory.get to return checkpoint with messages
+        mock_checkpoint = {
+            "channel_values": {
+                "messages": [
+                    {"role": "user", "content": "Question 1"},
+                    {"role": "assistant", "content": "Answer 1"},
+                    {"role": "user", "content": "Question 2"},
+                    {"role": "assistant", "content": "Answer 2"}
+                ]
+            }
+        }
+        mock_memory.get.return_value = mock_checkpoint
+        
+        messages = _retrieve_memory_impl(self.test_config)
+        
+        # Verify memory.get was called
+        mock_memory.get.assert_called_once_with(self.test_config)
+        
+        # Verify messages were retrieved
+        self.assertEqual(len(messages), 4)
+        self.assertEqual(messages[0]["content"], "Question 1")
+    
+    @patch('src.memory.memory')
+    def test_retrieve_memory_empty(self, mock_memory):
+        """Test retrieving memory when empty"""
+        # Mock memory.get to return None
+        mock_memory.get.return_value = None
+        
+        messages = _retrieve_memory_impl(self.test_config)
+        
+        # Should return empty list
+        self.assertEqual(messages, [])
+    
+    @patch('src.memory.memory')
+    def test_retrieve_memory_with_timer(self, mock_memory):
+        """Test retrieve_memory with timer"""
+        mock_memory.get.return_value = {}
+        mock_timer = Mock()
+        mock_timer.time_step = MagicMock()
+        mock_timer.time_step.return_value.__enter__ = Mock()
+        mock_timer.time_step.return_value.__exit__ = Mock()
+        
+        retrieve_memory(self.test_config, timer=mock_timer)
+        
+        # Verify timer was used
+        mock_timer.time_step.assert_called_once_with("memory_retrieval")
+    
+    @patch('src.memory.memory')
+    def test_checkpoint_structure(self, mock_memory):
+        """Test that checkpoint has correct structure"""
+        mock_memory.get.return_value = None
+        
+        _update_memory_impl(self.test_config, "Test", "Answer")
+        
+        call_args = mock_memory.put.call_args
+        checkpoint = call_args[0][1]
+        
+        # Verify checkpoint structure
+        self.assertIn("v", checkpoint)
+        self.assertIn("id", checkpoint)
+        self.assertIn("ts", checkpoint)
+        self.assertIn("channel_values", checkpoint)
+        self.assertIn("channel_versions", checkpoint)
+        self.assertIn("versions_seen", checkpoint)
+        self.assertEqual(checkpoint["v"], 1)
+
+
+if __name__ == '__main__':
+    unittest.main()

tests/test_response_generator.py

@@ -0,0 +1,200 @@
+"""
+Unit tests for response_generator module
+Tests LLM response generation functionality
+"""
+import unittest
+from unittest.mock import patch, Mock, MagicMock
+from src.response_generator import (
+    generate_xeno_response,
+    format_chat_history,
+    _generate_response_impl
+)
+
+
+class TestResponseGenerator(unittest.TestCase):
+    """Test cases for response_generator module"""
+    
+    def setUp(self):
+        """Set up test fixtures"""
+        self.context = """Knowledge Entry 1:
+Q: How do I create an account?
+A: Visit our website and click Sign Up.
+----------------------------------------"""
+        
+        self.question = "How can I create an account?"
+        
+        self.chat_history = [
+            {"role": "user", "content": "Hello"},
+            {"role": "assistant", "content": "Hi! How can I help you?"}
+        ]
+    
+    def test_format_chat_history(self):
+        """Test formatting chat history"""
+        formatted = format_chat_history(self.chat_history)
+        
+        # Check format
+        self.assertIn("User: Hello", formatted)
+        self.assertIn("Assistant: Hi! How can I help you?", formatted)
+        self.assertIn("\n", formatted)
+    
+    def test_format_chat_history_empty(self):
+        """Test formatting empty chat history"""
+        formatted = format_chat_history([])
+        self.assertEqual(formatted, "No previous conversation")
+    
+    def test_format_chat_history_single_message(self):
+        """Test formatting single message"""
+        history = [{"role": "user", "content": "Hello"}]
+        formatted = format_chat_history(history)
+        self.assertEqual(formatted, "User: Hello")
+    
+    def test_format_chat_history_missing_fields(self):
+        """Test formatting with missing fields"""
+        history = [
+            {"role": "user"},  # Missing content
+            {"content": "Test"}  # Missing role
+        ]
+        formatted = format_chat_history(history)
+        self.assertIn("User:", formatted)
+        self.assertIn("Unknown:", formatted)
+    
+    @patch('src.response_generator.genai.GenerativeModel')
+    def test_generate_response_impl(self, mock_model_class):
+        """Test internal response generation implementation"""
+        # Mock the model and response
+        mock_model = Mock()
+        mock_response = Mock()
+        mock_response.text = "You can create an account by visiting our website."
+        mock_model.generate_content.return_value = mock_response
+        mock_model_class.return_value = mock_model
+        
+        response = _generate_response_impl(
+            self.context, 
+            self.question, 
+            self.chat_history
+        )
+        
+        # Verify model was initialized with correct model name
+        mock_model_class.assert_called_once()
+        
+        # Verify generate_content was called
+        mock_model.generate_content.assert_called_once()
+        
+        # Check response
+        self.assertEqual(response, "You can create an account by visiting our website.")
+    
+    @patch('src.response_generator.genai.GenerativeModel')
+    def test_generate_response_with_empty_history(self, mock_model_class):
+        """Test generating response with empty history"""
+        mock_model = Mock()
+        mock_response = Mock()
+        mock_response.text = "Test response"
+        mock_model.generate_content.return_value = mock_response
+        mock_model_class.return_value = mock_model
+        
+        response = _generate_response_impl(
+            self.context, 
+            self.question, 
+            []
+        )
+        
+        # Verify it still works
+        self.assertEqual(response, "Test response")
+        
+        # Check that "None" was used for history in prompt
+        call_args = mock_model.generate_content.call_args
+        prompt = call_args[0][0]
+        self.assertIn("None", prompt)
+    
+    @patch('src.response_generator.genai.GenerativeModel')
+    def test_prompt_structure(self, mock_model_class):
+        """Test that prompt includes all necessary components"""
+        mock_model = Mock()
+        mock_response = Mock()
+        mock_response.text = "Test response"
+        mock_model.generate_content.return_value = mock_response
+        mock_model_class.return_value = mock_model
+        
+        _generate_response_impl(
+            self.context, 
+            self.question, 
+            self.chat_history
+        )
+        
+        # Get the prompt that was sent
+        call_args = mock_model.generate_content.call_args
+        prompt = call_args[0][0]
+        
+        # Verify prompt structure
+        self.assertIn("HISTORY", prompt)
+        self.assertIn("CONTEXT", prompt)
+        self.assertIn("QUESTION", prompt)
+        self.assertIn(self.context, prompt)
+        self.assertIn(self.question, prompt)
+    
+    @patch('src.response_generator.genai.GenerativeModel')
+    def test_generate_xeno_response_with_timer(self, mock_model_class):
+        """Test generate_xeno_response with timer"""
+        mock_model = Mock()
+        mock_response = Mock()
+        mock_response.text = "Test response"
+        mock_model.generate_content.return_value = mock_response
+        mock_model_class.return_value = mock_model
+        
+        mock_timer = Mock()
+        mock_timer.time_step = MagicMock()
+        mock_timer.time_step.return_value.__enter__ = Mock()
+        mock_timer.time_step.return_value.__exit__ = Mock()
+        
+        response = generate_xeno_response(
+            self.context, 
+            self.question, 
+            self.chat_history, 
+            timer=mock_timer
+        )
+        
+        # Verify timer was used
+        mock_timer.time_step.assert_called_once_with("llm_generation")
+        
+        # Verify response
+        self.assertEqual(response, "Test response")
+    
+    @patch('src.response_generator.genai.GenerativeModel')
+    def test_response_text_stripping(self, mock_model_class):
+        """Test that response text is stripped of whitespace"""
+        mock_model = Mock()
+        mock_response = Mock()
+        mock_response.text = "  Test response with spaces  \n"
+        mock_model.generate_content.return_value = mock_response
+        mock_model_class.return_value = mock_model
+        
+        response = _generate_response_impl(
+            self.context, 
+            self.question, 
+            []
+        )
+        
+        # Should be stripped
+        self.assertEqual(response, "Test response with spaces")
+    
+    @patch('src.response_generator.genai.GenerativeModel')
+    def test_system_prompt_inclusion(self, mock_model_class):
+        """Test that system prompt is included in generated prompt"""
+        mock_model = Mock()
+        mock_response = Mock()
+        mock_response.text = "Test"
+        mock_model.generate_content.return_value = mock_response
+        mock_model_class.return_value = mock_model
+        
+        _generate_response_impl(self.context, self.question, [])
+        
+        # Get the prompt
+        call_args = mock_model.generate_content.call_args
+        prompt = call_args[0][0]
+        
+        # Should contain system prompt text
+        self.assertIn("XENO Support Assistant", prompt)
+
+
+if __name__ == '__main__':
+    unittest.main()

tests/test_utils.py

@@ -0,0 +1,108 @@
+"""
+Unit tests for utils module
+Tests the PipelineTimer class
+"""
+import unittest
+import time
+from src.utils import PipelineTimer
+
+
+class TestPipelineTimer(unittest.TestCase):
+    """Test cases for PipelineTimer class"""
+    
+    def setUp(self):
+        """Set up test fixtures"""
+        self.timer = PipelineTimer()
+    
+    def test_initialization(self):
+        """Test timer initialization"""
+        self.assertIsNotNone(self.timer.start_time)
+        self.assertEqual(self.timer.step_times, {})
+        self.assertIsNone(self.timer.step_start)
+        self.assertIsNone(self.timer.current_step)
+    
+    def test_reset(self):
+        """Test timer reset functionality"""
+        # Add some data
+        self.timer.step_times = {'test': 100}
+        self.timer.current_step = 'test'
+        
+        # Reset
+        self.timer.reset()
+        
+        # Verify reset
+        self.assertEqual(self.timer.step_times, {})
+        self.assertIsNone(self.timer.current_step)
+    
+    def test_time_step_context_manager(self):
+        """Test timing a step using context manager"""
+        with self.timer.time_step('test_step'):
+            time.sleep(0.1)  # Sleep for 100ms
+        
+        # Check that step was timed
+        self.assertIn('test_step', self.timer.step_times)
+        # Should be approximately 100ms (allowing some variance)
+        self.assertGreater(self.timer.step_times['test_step'], 90)
+        self.assertLess(self.timer.step_times['test_step'], 150)
+    
+    def test_multiple_steps(self):
+        """Test timing multiple steps"""
+        with self.timer.time_step('step1'):
+            time.sleep(0.05)
+        
+        with self.timer.time_step('step2'):
+            time.sleep(0.05)
+        
+        # Both steps should be recorded
+        self.assertIn('step1', self.timer.step_times)
+        self.assertIn('step2', self.timer.step_times)
+        self.assertEqual(len(self.timer.step_times), 2)
+    
+    def test_get_total_time(self):
+        """Test getting total elapsed time"""
+        time.sleep(0.1)
+        total_time = self.timer.get_total_time()
+        
+        # Should be at least 100ms
+        self.assertGreater(total_time, 90)
+    
+    def test_get_timing_summary(self):
+        """Test getting timing summary"""
+        with self.timer.time_step('step1'):
+            time.sleep(0.05)
+        
+        summary = self.timer.get_timing_summary()
+        
+        # Check summary structure
+        self.assertIn('total_time_ms', summary)
+        self.assertIn('step_times', summary)
+        self.assertIn('timestamp', summary)
+        self.assertIn('step1', summary['step_times'])
+    
+    def test_current_step_tracking(self):
+        """Test that current_step is tracked correctly"""
+        self.assertIsNone(self.timer.current_step)
+        
+        with self.timer.time_step('test_step'):
+            # During execution, current_step should be set
+            self.assertEqual(self.timer.current_step, 'test_step')
+        
+        # After execution, current_step should be None
+        self.assertIsNone(self.timer.current_step)
+    
+    def test_exception_handling_in_timer(self):
+        """Test that timer handles exceptions properly"""
+        try:
+            with self.timer.time_step('error_step'):
+                raise ValueError("Test error")
+        except ValueError:
+            pass
+        
+        # Step should still be recorded even if exception occurred
+        self.assertIn('error_step', self.timer.step_times)
+        # current_step should be None after context manager exits
+        self.assertIsNone(self.timer.current_step)
+
+
+if __name__ == '__main__':
+    unittest.main()

tests/test_vector_store.py

@@ -0,0 +1,243 @@
+"""
+Unit tests for vector_store module
+Tests ChromaDB vector store operations
+"""
+import unittest
+import numpy as np
+import torch
+from unittest.mock import patch, Mock, MagicMock
+from src.vector_store import (
+    generate_embeddings,
+    calculate_similarity,
+    process_context,
+    _generate_embeddings_impl,
+    _calculate_similarity_impl,
+    _process_context_impl
+)
+
+
+class TestVectorStore(unittest.TestCase):
+    """Test cases for vector_store module"""
+    
+    def setUp(self):
+        """Set up test fixtures"""
+        # Mock document
+        self.mock_doc = Mock()
+        self.mock_doc.page_content = "Test document content"
+        self.mock_doc.metadata = {
+            'id': 'KB001',
+            'question': 'Test question?',
+            'content': 'Test answer.',
+            'section': 'Test'
+        }
+        
+        self.mock_documents = [self.mock_doc]
+    
+    @patch('src.vector_store.genai.embed_content')
+    def test_generate_embeddings_impl(self, mock_embed):
+        """Test internal embedding generation implementation"""
+        # Mock embeddings
+        mock_embed.side_effect = [
+            {'embedding': [0.1, 0.2, 0.3]},  # Query embedding
+            {'embedding': [0.2, 0.3, 0.4]}   # Doc embedding
+        ]
+        
+        query = "Test query"
+        query_emb, doc_embs = _generate_embeddings_impl(query, self.mock_documents)
+        
+        # Verify embed_content was called correctly
+        self.assertEqual(mock_embed.call_count, 2)
+        
+        # Check query embedding call
+        first_call = mock_embed.call_args_list[0]
+        self.assertEqual(first_call[1]['content'], query)
+        self.assertEqual(first_call[1]['task_type'], 'retrieval_query')
+        
+        # Check doc embedding call
+        second_call = mock_embed.call_args_list[1]
+        self.assertEqual(second_call[1]['content'], self.mock_doc.page_content)
+        self.assertEqual(second_call[1]['task_type'], 'retrieval_document')
+        
+        # Verify embeddings
+        self.assertEqual(query_emb, [0.1, 0.2, 0.3])
+        self.assertEqual(len(doc_embs), 1)
+        self.assertEqual(doc_embs[0], [0.2, 0.3, 0.4])
+    
+    @patch('src.vector_store.genai.embed_content')
+    def test_generate_embeddings_with_timer(self, mock_embed):
+        """Test embedding generation with timer"""
+        mock_embed.side_effect = [
+            {'embedding': [0.1, 0.2, 0.3]},
+            {'embedding': [0.2, 0.3, 0.4]}
+        ]
+        
+        mock_timer = Mock()
+        mock_timer.time_step = MagicMock()
+        mock_timer.time_step.return_value.__enter__ = Mock()
+        mock_timer.time_step.return_value.__exit__ = Mock()
+        
+        generate_embeddings("Test", self.mock_documents, timer=mock_timer)
+        
+        # Verify timer was used
+        mock_timer.time_step.assert_called_once_with("embedding_generation")
+    
+    @patch('src.vector_store.genai.embed_content')
+    def test_generate_embeddings_multiple_docs(self, mock_embed):
+        """Test embedding generation with multiple documents"""
+        # Create multiple mock documents
+        mock_doc2 = Mock()
+        mock_doc2.page_content = "Second document"
+        docs = [self.mock_doc, mock_doc2]
+        
+        # Mock embeddings
+        mock_embed.side_effect = [
+            {'embedding': [0.1, 0.2, 0.3]},  # Query
+            {'embedding': [0.2, 0.3, 0.4]},  # Doc 1
+            {'embedding': [0.3, 0.4, 0.5]}   # Doc 2
+        ]
+        
+        query_emb, doc_embs = _generate_embeddings_impl("Test", docs)
+        
+        # Should have 2 doc embeddings
+        self.assertEqual(len(doc_embs), 2)
+        self.assertEqual(mock_embed.call_count, 3)
+    
+    def test_calculate_similarity_impl(self):
+        """Test internal similarity calculation implementation"""
+        query_embedding = [1.0, 0.0, 0.0]
+        doc_embeddings = [
+            [1.0, 0.0, 0.0],  # Same as query - score should be ~1.0
+            [0.0, 1.0, 0.0],  # Orthogonal - score should be ~0.0
+            [0.5, 0.5, 0.0]   # Partial similarity
+        ]
+        
+        scores = _calculate_similarity_impl(query_embedding, doc_embeddings)
+        
+        # Check scores
+        self.assertEqual(len(scores), 3)
+        self.assertAlmostEqual(scores[0], 1.0, places=5)
+        self.assertAlmostEqual(scores[1], 0.0, places=5)
+        self.assertGreater(scores[2], 0.0)
+        self.assertLess(scores[2], 1.0)
+    
+    def test_calculate_similarity_with_timer(self):
+        """Test similarity calculation with timer"""
+        mock_timer = Mock()
+        mock_timer.time_step = MagicMock()
+        mock_timer.time_step.return_value.__enter__ = Mock()
+        mock_timer.time_step.return_value.__exit__ = Mock()
+        
+        query_emb = [1.0, 0.0, 0.0]
+        doc_embs = [[1.0, 0.0, 0.0]]
+        
+        calculate_similarity(query_emb, doc_embs, timer=mock_timer)
+        
+        # Verify timer was used
+        mock_timer.time_step.assert_called_once_with("similarity_calculation")
+    
+    def test_process_context_impl(self):
+        """Test internal context processing implementation"""
+        # Create mock results with metadata
+        results = []
+        for i in range(3):
+            mock_result = Mock()
+            mock_result.metadata = {
+                'id': f'KB00{i+1}',
+                'question': f'Question {i+1}?',
+                'content': f'Answer {i+1}.'
+            }
+            results.append(mock_result)
+        
+        # Cosine scores (sorted: 0.9, 0.7, 0.5)
+        cosine_scores = [0.7, 0.5, 0.9]
+        
+        context, source_ids, knowledge_pairs = _process_context_impl(
+            results, cosine_scores, max_results=2
+        )
+        
+        # Should return top 2 results
+        self.assertEqual(len(source_ids), 2)
+        self.assertEqual(len(knowledge_pairs), 2)
+        
+        # Check that highest score (0.9, index 2) is first
+        self.assertEqual(source_ids[0], 'KB003')
+        self.assertEqual(knowledge_pairs[0][0], 'Question 3?')
+        
+        # Check formatted context
+        self.assertIn("Knowledge Entry 1:", context)
+        self.assertIn("Knowledge Entry 2:", context)
+        self.assertIn("Q: Question 3?", context)
+        self.assertIn("A: Answer 3.", context)
+    
+    def test_process_context_with_timer(self):
+        """Test context processing with timer"""
+        mock_result = Mock()
+        mock_result.metadata = {'id': 'KB001', 'question': 'Q?', 'content': 'A.'}
+        
+        mock_timer = Mock()
+        mock_timer.time_step = MagicMock()
+        mock_timer.time_step.return_value.__enter__ = Mock()
+        mock_timer.time_step.return_value.__exit__ = Mock()
+        
+        process_context([mock_result], [0.9], timer=mock_timer)
+        
+        # Verify timer was used
+        mock_timer.time_step.assert_called_once_with("context_processing")
+    
+    def test_process_context_max_results(self):
+        """Test that max_results parameter limits output"""
+        # Create 5 mock results
+        results = []
+        for i in range(5):
+            mock_result = Mock()
+            mock_result.metadata = {
+                'id': f'KB00{i}',
+                'question': f'Q{i}?',
+                'content': f'A{i}.'
+            }
+            results.append(mock_result)
+        
+        scores = [0.9, 0.8, 0.7, 0.6, 0.5]
+        
+        # Request only 3 results
+        context, source_ids, knowledge_pairs = _process_context_impl(
+            results, scores, max_results=3
+        )
+        
+        # Should only return 3
+        self.assertEqual(len(source_ids), 3)
+        self.assertEqual(len(knowledge_pairs), 3)
+    
+    def test_process_context_formatting(self):
+        """Test context formatting details"""
+        mock_result = Mock()
+        mock_result.metadata = {
+            'id': 'KB001',
+            'question': 'Test question?',
+            'content': 'Test answer.'
+        }
+        
+        context, _, _ = _process_context_impl([mock_result], [0.9], max_results=1)
+        
+        # Check formatting
+        self.assertIn("Knowledge Entry 1:", context)
+        self.assertIn("Q: Test question?", context)
+        self.assertIn("A: Test answer.", context)
+        self.assertIn("-" * 40, context)
+    
+    def test_process_context_missing_metadata(self):
+        """Test context processing with missing metadata fields"""
+        mock_result = Mock()
+        mock_result.metadata = {}  # No metadata
+        
+        context, source_ids, knowledge_pairs = _process_context_impl(
+            [mock_result], [0.9], max_results=1
+        )
+        
+        # Should handle missing fields with N/A
+        self.assertIn("N/A", context)
+        self.assertEqual(source_ids[0], "N/A")
+
+
+if __name__ == '__main__':
+    unittest.main()