Baseline

Author: seachellemz

.gitignore

@@ -0,0 +1,7 @@
+*.pyc
+.DS_Store
+backup
+chroma_boardgames/
+env/
+.env
+__pycache__

README.md

@@ -0,0 +1,75 @@
+# Simple Local RAG
+
+A basic local LLM RAG Chatbot With LangChain that exposes itself via REST endpoints.
+
+## Setup
+
+Clone this repository and create a clean python v3.10 virtual environment and activate it.
+
+#### Dependencies
+
+The following is a high-level list of components used to run this local RAG:
+
+- langchain
+- streamlit
+- streamlit-chat
+- pypdf
+- chromadb
+- fastembed
+
+```
+pip install -r requirements.txt
+```
+
+#### Setup up Ollama
+
+This depends on the Ollama platform to run the LLM locally. The setup is straightforward. First, visit ollama.ai and download the app appropriate for your operating system.
+
+Next open your terminal and execute the following command to pull the latest Mistral model.
+
+```
+ollama pull llama3
+```
+
+#### Configuration
+
+Create a `.env` file in the root directory and add the following environment variables:
+
+```.env
+
+CHROMA_PATH=chroma_boardgames
+DATA_PATH_BG=data_boardgames
+```
+
+#### Build the Vector Store
+
+The `populate_database.py` loads any PDF files it finds in the `DATA_PATH_BG` folder. The repository currently includes a couple of example board game instruction manuals to seed a Chroma Vector store. The module reads the folder and loads each of the PDF's it into vector storage in two steps: first, it splits the document into smaller chunks to accommodate the token limit of the LLM; second, it vectorizes these chunks using FastEmbeddings and stores them into Chroma. It will generate a chunk ID that will indicate which PDF file, page number and chunk number of the embedding. This allows us to analyze how the model is producing a response, but also allows us to incrementally add new data to the database without have to fully reload it. Run the database load:
+
+` python -m populate_database.py`
+
+If you need to clear the database for any reason, run:
+
+`python -m reset_database.py`
+
+The above command will remove the chroma database. If you need to recreate it, simply rerun `populate_database.py`
+
+## Running the RAG
+
+The instruction manuals for both Monopoly and Ticket To Ride have been loaded into the Chroma DB. Ask the RAG questions about these two board games and see how well it does answering your questions. The RAG can be invoked using the following command with the sample question:
+
+```
+python query_data.py  How do I build a hotel in monopoly?
+```
+
+Here are some additional questions you can try:
+
+- How much total money does a player start with in Monopoly? (Answer with the number only)
+- How many points does the longest continuous train get in Ticket to Ride? (Answer with the number only)
+
+You can also browse the instruction manuals that are in the `./data_boardgames` folder to come up with your own questions.
+
+## Running the test cases
+
+```
+pytest test_rag.py
+```

api.py

@@ -0,0 +1,100 @@
+import logging
+from query_data import query_rag
+from fastapi import FastAPI, Request, HTTPException, status
+from fastapi.exceptions import RequestValidationError
+from fastapi.responses import JSONResponse
+from pydantic import BaseModel
+from models.rag_query_model import QueryInput, QueryOutput
+from utils.async_utils import async_retry
+
+
+class Message(BaseModel):
+    """ Message class defined in Pydantic """
+    channel: str
+    author: str
+    text: str
+
+
+app = FastAPI(
+    title="PDF Document Chatbot",
+    description="Endpoints for various PDF documents",
+)
+
+channel_list = ["general", "dev", "marketing"]
+message_map = {}
+
+
+@app.exception_handler(RequestValidationError)
+async def validation_exception_handler(request: Request, exc: RequestValidationError):
+    exc_str = f'{exc}'.replace('\n', ' ').replace('   ', ' ')
+    logging.error(f"{request}: {exc_str}")
+    content = {'status_code': 10422, 'message': exc_str, 'data': None}
+    return JSONResponse(content=content, status_code=status.HTTP_422_UNPROCESSABLE_ENTITY)
+
+
+@async_retry(max_retries=10, delay=1)
+async def invoke_agent_with_retry(query: str):
+    """
+    Retry the agent if a tool fails to run. This can help when there
+    are intermittent connection issues to external APIs.
+    """
+
+    return await query_rag({"input": query})
+
+
+@app.get("/")
+async def get_status():
+    return {"status": "running"}
+
+
+@app.post("/post_message", status_code=status.HTTP_201_CREATED)
+def post_message(message: Message):
+    """Post a new message to the specified channel."""
+    channel = message.channel
+    if channel in channel_list:
+        # message_map[channel].append(message)
+        return message
+    else:
+        raise HTTPException(status_code=404, detail="channel not found")
+
+
+@app.post("/rag-query")
+async def query_rag_api(query: QueryInput):
+    print(f"api.py - API Request Data: {query}")
+    query_response = query_rag({"input": query})
+    print(query_response)
+
+    # query: str
+    # response: str
+    # sources: list[str]
+    query_response2 = {
+        "query": query, "response": query_response["response"], "sources": query_response["sources"]}
+    print(f"Query Response2:  {query_response2}")
+
+    # query_response["intermediate_steps"] = [
+    #    str(s) for s in query_response["intermediate_steps"]
+    # ]
+
+    return query_response2
+
+
+@app.post("/rag-query2")
+async def query_rag_api2(
+    query: QueryInput,
+) -> QueryOutput:
+    query_response = query_rag({"input": query})
+    print(query_response)
+
+    # query: str
+    # response: str
+    # sources: list[str]
+    query_text = query["query"]
+    query_response2 = {
+        "query": query_text, "response": query_response["response"], "sources": query_response["sources"]}
+    print(f"Query Response2:  {query_response2}")
+
+    # query_response["intermediate_steps"] = [
+    #    str(s) for s in query_response["intermediate_steps"]
+    # ]
+
+    return query_response2