SAGE: Personal RAG-enabled Information Retrieval Utility

SAGE is a complete RAG (Retrieval-Augmented Generation) management and interrogation system built with FastAPI and ChromaDB. It provides a powerful API for document management, vector storage, and LLM-powered chat functionality with support for multiple LLM providers.

Features

• Document Management: Upload, process, and manage documents in vector collections
• Document Tagging: Add tags to documents for better organization and filtering
• Vector Database: ChromaDB integration for efficient semantic search
• LLM Integration: Support for multiple LLM providers:
  • Ollama with local or remote models (including nomic-embed-text)
  • Anthropic with Claude models (including haiku)
  • OpenAI with GPT models (including 4o)
• Conversational UI: Interactive chat interface with conversation history and context awareness
• WebSocket Chat: Real-time, streaming chat interface with RAG capabilities
• REST API: Comprehensive REST API for all operations
• Document Processing: Support for various document types (PDF, DOCX, CSV, text)

Installation

Prerequisites

• Python 3.8+
• [Optional] Virtual environment (recommended)

Setup

1. Clone the repository:

   jj git clone --colocate https://github.com/grepory/sage.git
   cd sage

2. Create and activate a virtual environment (optional but recommended):

   python -m venv venv
   source venv/bin/activate # On Windows: venv\Scripts\activate

3. Install dependencies:

   pip install -r requirements.txt

4. Copy the .env.example file to .env in the project root and update it with your LLM provider configuration:

   cp .env.example .env
   # Then edit the .env file with your preferred text editor

   The .env.example file contains all the necessary environment variables with default values and explanatory comments. You only need to configure the providers you plan to use.

Running the Application

Start the server with:

python run_app.py

This will:

• Start the server on the configured port (default 8000)
• Display clear URLs for accessing the application
• Automatically open your default web browser to the frontend

Alternatively, you can start the server manually with:

uvicorn app.main:app --reload --host 0.0.0.0 --port 8000

The API will be available at http://localhost:8000, and the API documentation at http://localhost:8000/docs.

Configuration

Sage can be configured through environment variables in the .env file:

Server Configuration

• PORT=8000 - Port for the FastAPI server
• ROOT_PATH=/sage - Root path prefix for reverse proxy setup (e.g., for hosting at http://example.com/sage/)

LLM Provider Configuration

Configure your preferred LLM providers in the .env file:

API Documentation

Collections API

Collections are used to organize documents in the vector database.

Create a Collection

POST /api/v1/collections/Content-Type: application/json
{
"name": "my_collection",
"description": "My first collection"
}

List Collections

GET /api/v1/collections/

Get Collection

GET /api/v1/collections/{collection_name}

Delete Collection

DELETE /api/v1/collections/{collection_name}

Documents API

Documents are the core content stored in the vector database.

Upload a Document

POST /api/v1/documents/uploadContent-Type: multipart/form-datafile: [file]collection_name: my_collectiontags: personal,house,importantadditional_metadata: {"source": "website", "author": "John Doe"}

The tags parameter is optional and should be a comma-separated list of tags to associate with the document.

Add Text Directly

POST /api/v1/documents/textContent-Type: application/json
{
"text": "This is some text to add to the vector database.",
"collection_name": "my_collection",
"tags": ["personal", "notes", "important"],
"metadata": {"source": "direct_input", "author": "Jane Smith"}
}

The tags parameter is optional and should be an array of strings representing tags to associate with the document.

Get Document

GET /api/v1/documents/{collection_name}/{document_id}

Delete Document

DELETE /api/v1/documents/{collection_name}/{document_id}

Query Documents

POST /api/v1/documents/queryContent-Type: application/json
{
"collection_name": "my_collection",
"query_text": "What is retrieval-augmented generation?",
"n_results": 5,
"where": {
"tags": {"$in": ["personal", "important"]}
}
}

The where parameter is optional and can be used to filter documents by metadata fields, including tags. In the example above, the query will only return documents that have either the "personal" or "important" tag. You can use other operators like $eq for exact match or $contains for substring match.

Chat API

The Chat API provides both REST and WebSocket interfaces for interacting with the LLM.

Model Selection

You can specify which LLM provider and model to use by setting the model parameter in your requests. The format is:

provider:model

For example:

• ollama:llama2 - Use Ollama with the llama2 model
• anthropic:claude-3-haiku-20240307 - Use Anthropic with the Claude 3 Haiku model
• openai:gpt-4o - Use OpenAI with the GPT-4o model

If you don't specify a provider, the default provider from your configuration will be used. If you don't specify a model, the default model for the selected provider will be used.

REST Chat Endpoint

POST /api/v1/chat/Content-Type: application/json
{
"collection_name": "my_collection",
"query": "What is retrieval-augmented generation?",
"history": [
{"role": "user", "content": "Hello"},
{"role": "assistant", "content": "Hi there! How can I help you?"}
],
"model": "llama2"
}

WebSocket Chat

Connect to the WebSocket endpoint at /api/v1/chat/ws and send JSON messages with the following format:

{
"type": "chat",
"collection_name": "my_collection",
"query": "What is retrieval-augmented generation?",
"history": [
{"role": "user", "content": "Hello"},
{"role": "assistant", "content": "Hi there! How can I help you?"}
],
"model": "llama2"
}

The server will respond with messages in the following formats:

1. Start message:

{
"type": "start",
"content": "Generating response..."
}

2. Token messages (streaming):

{
"type": "token",
"content": "piece of text"
}

3. Complete message:

{
"type": "complete",
"content": {
"answer": "The complete answer...",
"sources": [
{
"id": "document_id",
"text": "Snippet of the source document...",
"metadata": {"source": "document.pdf", "page": 1}
}
],
"history": [
{"role": "user", "content": "What is RAG?"},
{"role": "assistant", "content": "The complete answer..."}
]
}
}

4. Error message:

{
"type": "error",
"content": "Error message"
}

Web Interface

Sage provides a responsive web interface that works seamlessly on desktop and mobile devices. The interface includes:

Document Upload

• Upload files using the upload button
• Add tags to organize your documents
• View upload progress and status

Conversational Chat Interface

The chat interface allows you to have natural conversations with the AI about your documents:

• Conversation History: The system maintains the full conversation history, allowing for contextual follow-up questions
• Conversation Sidebar: Browse and manage your chat history with a collapsible sidebar
• Enter Key Support: Press Enter to send your message (use Shift+Enter for new lines)
• Visual Distinction: User and assistant messages are visually distinct for easy reading
• Source Citations: The system displays the source documents used to generate responses
• Mobile Responsive: Optimized interface for mobile devices with touch-friendly controls
• Tag-based Filtering: Filter conversations and documents by tags with advanced search capabilities

To use the conversational interface:

1. Select a a set of tags from the dropdown (or choose to search all tags)
2. (Optional) Select a specific LLM model and provider
3. Type your question in the message box
4. Press Enter or click "Send Message"
5. View the AI's response and the sources used
6. Continue the conversation with follow-up questions

The system will maintain context between questions, allowing for a more natural conversation flow. On mobile devices, the conversation sidebar can be toggled using the sidebar button in the top navigation.

Examples

Python Client Example

Here's a simple Python client for interacting with the API:

importrequestsimportwebsocketsimportjsonimportasyncioBASE_URL="http://localhost:8000/api/v1"# Create a collectiondefcreate_collection(name, description=None):
response=requests.post(
f"{BASE_URL}/collections/",
json={"name": name, "description": description}
)
returnresponse.json()
# Upload a documentdefupload_document(file_path, collection_name, tags=None, metadata=None):
files= {"file": open(file_path, "rb")}
data= {"collection_name": collection_name}
iftags:
data["tags"] =",".join(tags)
ifmetadata:
data["additional_metadata"] =json.dumps(metadata)
response=requests.post(
f"{BASE_URL}/documents/upload",
files=files,
data=data
)
returnresponse.json()
# Query documentsdefquery_documents(collection_name, query_text, n_results=5, tags=None):
query_data= {
"collection_name": collection_name,
"query_text": query_text,
"n_results": n_results
}
# Add tag filtering if specifiediftags:
query_data["where"] = {"tags": {"$in": tags}}
response=requests.post(
f"{BASE_URL}/documents/query",
json=query_data
)
returnresponse.json()
# WebSocket chat clientasyncdefchat_websocket(query, tags=None, history=None, model=None, include_untagged=True):
uri=f"ws://localhost:8000/api/v1/chat/ws"asyncwithwebsockets.connect(uri) aswebsocket:
# Send chat messagemessage= {
"type": "chat",
"query": query,
"history": historyor [],
"include_untagged": include_untagged
}
iftags:
message["tags"] =tagsifmodel:
message["model"] =modelawaitwebsocket.send(json.dumps(message))
# Process responsesfull_response=""whileTrue:
response=awaitwebsocket.recv()
data=json.loads(response)
ifdata["type"] =="token":
# Print token as it arrivesprint(data["content"], end="", flush=True)
full_response+=data["content"]
elifdata["type"] =="complete":
# Chat completeprint("\n\nSources:")
fori, sourceinenumerate(data["content"]["sources"]):
print(f"{i+1}. {source['text']} (ID: {source['id']})")
returndata["content"]
elifdata["type"] =="error":
# Error occurredprint(f"\nError: {data['content']}")
returnNone# Example usageif__name__=="__main__":
# Create a collectioncreate_collection("my_docs", "My documents collection")
# Upload a document with tagsupload_document(
"path/to/document.pdf", "my_docs", tags=["house", "important", "insurance"],
metadata={"source": "local", "author": "Insurance Company"}
)
# Upload another document with different tagsupload_document(
"path/to/another_document.pdf", "my_docs", tags=["personal", "important", "legal"],
metadata={"source": "local", "author": "Law Firm"}
)
# Query documents without tag filteringresults=query_documents("my_docs", "What is RAG?")
# Query documents with tag filtering (only "house" tagged documents)house_results=query_documents("my_docs", "What is RAG?", tags=["house"])
# Query documents with multiple tag filtering (documents tagged as either "important" or "personal")important_results=query_documents("my_docs", "What is RAG?", tags=["important", "personal"])
# Chat with WebSocket - Single queryasyncio.run(chat_websocket("Explain RAG in simple terms."))
# Chat with WebSocket - Query with tag filteringasyncio.run(chat_websocket(
"What can you tell me about my house documents?", tags=["house", "important"]
))
# Chat with WebSocket - Conversation with historyasyncdefconversation_example():
# First queryprint("\n--- Starting conversation ---")
response=awaitchat_websocket("Explain RAG in simple terms.")
# Get the updated history from the responseconversation_history=response["history"]
# Follow-up query using the conversation history with tag filteringprint("\n--- Follow-up question with tag filtering ---")
response=awaitchat_websocket(
"What are the main benefits compared to traditional approaches?",
tags=["important", "personal"],
history=conversation_history
)
# Continue the conversation with another follow-upconversation_history=response["history"]
print("\n--- Another follow-up question ---")
awaitchat_websocket(
"Can you give me a simple code example?",
history=conversation_history
)
# Run the conversation exampleasyncio.run(conversation_example())

License

MIT

Contributing

Contributions are welcome! Please feel free to submit a Pull Request.