feat: Use YAML config for HuggingFace token - Add confg.yaml support for HuggingFace authentication - Update local_rag_agent to load token from config - Add config_example.yaml template - Update documentation for YAML-based configuration

Author: jasperan

agentic_rag/README.md

@@ -25,19 +25,15 @@ The system has the following features:
 
    The system uses Mistral-7B by default, which requires authentication with HuggingFace:
 
-   a. Create a HuggingFace account at https://huggingface.co/join
+   a. Create a HuggingFace account [here](https://huggingface.co/join)
 
-   b. Accept the Mistral-7B model terms at https://huggingface.co/mistralai/Mistral-7B-Instruct-v0.2
+   b. Accept the Mistral-7B model terms & conditions [here](https://huggingface.co/mistralai/Mistral-7B-Instruct-v0.2)
 
-   c. Create an access token at https://huggingface.co/settings/tokens
+   c. Create an access token [here](https://huggingface.co/settings/tokens)
 
-   d. Login using the token:
-   ```bash
-   huggingface-cli login
-   # Or set the token as an environment variable:
-   export HUGGING_FACE_HUB_TOKEN=your_token_here
-   # On Windows:
-   set HUGGING_FACE_HUB_TOKEN=your_token_here
+   d. Create a `config.yaml` file (you can copy from `config_example.yaml`):
+   ```yaml
+   HUGGING_FACE_HUB_TOKEN: your_token_here
    ```
 
 3. (Optional) If you want to use the OpenAI-based agent instead of the default local model, create a `.env` file with your OpenAI API key:
@@ -46,9 +42,7 @@ The system has the following features:
    OPENAI_API_KEY=your-api-key-here
    ```
 
-   If no API key is provided, the system will automatically use the local Mistral-7B model for text generation.
-
-4. The system will automatically download and use `Mistral-7B-Instruct-v0.2` for text generation when using the local model. No additional configuration is needed.
+4. If no API key is provided, the system will automatically download and use `Mistral-7B-Instruct-v0.2` for text generation when using the local model. No additional configuration is needed.
 
 ## 1. Getting Started
 

agentic_rag/config_example.yaml

@@ -0,0 +1 @@
+HUGGING_FACE_HUB_TOKEN: your_token_here

agentic_rag/local_rag_agent.py

@@ -4,6 +4,8 @@
 from pydantic import BaseModel, Field
 from store import VectorStore
 import argparse
+import yaml
+import os
 
 class QueryAnalysis(BaseModel):
     """Pydantic model for query analysis output"""
@@ -22,14 +24,25 @@ def __init__(self, vector_store: VectorStore, model_name: str = "mistralai/Mistr
         """Initialize local RAG agent with vector store and local LLM"""
         self.vector_store = vector_store
 
+        # Load HuggingFace token from config
+        try:
+            with open('config.yaml', 'r') as f:
+                config = yaml.safe_load(f)
+            token = config.get('HUGGING_FACE_HUB_TOKEN')
+            if not token:
+                raise ValueError("HUGGING_FACE_HUB_TOKEN not found in config.yaml")
+        except Exception as e:
+            raise Exception(f"Failed to load HuggingFace token from config.yaml: {str(e)}")
+        
         # Load model and tokenizer
         print("\nLoading model and tokenizer...")
         self.model = AutoModelForCausalLM.from_pretrained(
             model_name,
             torch_dtype=torch.float16,
-            device_map="auto"
+            device_map="auto",
+            token=token
         )
-        self.tokenizer = AutoTokenizer.from_pretrained(model_name)
+        self.tokenizer = AutoTokenizer.from_pretrained(model_name, token=token)
 
         # Create text generation pipeline
         self.pipeline = pipeline(

agentic_rag/requirements.txt

@@ -9,4 +9,5 @@ uvicorn
 python-multipart
 transformers
 torch
-accelerate
+accelerate
+pyyaml