rag_system package

RAG System for Stargazer Bot v3.

A file-based Retrieval-Augmented Generation system using: - Gemini API for embeddings (google/gemini-embedding-001) via shared key pool - ChromaDB for vector storage - Full file retrieval (not chunks) - Per-channel auto-search with context injection

class rag_system.OpenRouterEmbeddings(api_key=None, model='google/gemini-embedding-001', dimensions=None, timeout=30.0, gemini_api_key=None, gemini_only=True)

Bases: object

Async embeddings client using Gemini API via the shared key pool.

Parameters:

• api_key (str | None)

• model (str)

• dimensions (int | None)

• timeout (float)

• gemini_api_key (str | None)

• gemini_only (bool)

DEFAULT_MODEL = 'google/gemini-embedding-001'

MAX_BATCH_SIZE = 50

MAX_BATCH_CHARS = 50000

__init__(api_key=None, model='google/gemini-embedding-001', dimensions=None, timeout=30.0, gemini_api_key=None, gemini_only=True)

Initialize the instance.

Parameters:

• api_key (Optional[str]) – Unused; kept for backward compatibility.

• model (str) – The model value.

• dimensions (Optional[int]) – The dimensions value.

• timeout (float) – Maximum wait time in seconds.

• gemini_api_key (Optional[str]) – Unused; pool is used instead.

• gemini_only (bool) – Always True; embeddings use Gemini API only.

async embed_text(text)

Embed text.

Parameters:

text (str) – Text content.

Returns:

The result.

Return type:

ndarray

async embed_texts(texts)

Embed texts.

Parameters:

texts (List[str]) – The texts value.

Returns:

The result.

Return type:

List[ndarray]

async embed_text_for_search(text, task_type='QUESTION_ANSWERING')

Embed a single text using the Gemini API only, with a task type.

Intended for pre-computing a query embedding before passing it to FileRAGManager.search(query_embedding=...). Retries on transient errors with exponential back-off.

Return type:

List[float]

Parameters:

• text (str)

• task_type (str)

async close()

Close.

async __aenter__()

Internal helper: aenter .

async __aexit__(exc_type, exc_val, exc_tb)

Internal helper: aexit .

Parameters:

• exc_type – The exc type value.

• exc_val – The exc val value.

• exc_tb – The exc tb value.

class rag_system.SyncOpenRouterEmbeddings(api_key=None, model='google/gemini-embedding-001', dimensions=None, timeout=30.0, gemini_api_key=None, gemini_only=True, document_task_type=None, query_task_type=None)

Bases: object

Synchronous wrapper used by ChromaDB’s embedding function interface.

Uses Gemini API via the shared key pool. Batches are dispatched concurrently via a ThreadPoolExecutor when there are multiple batches.

Parameters:

• api_key (str | None)

• model (str)

• dimensions (int | None)

• timeout (float)

• gemini_api_key (str | None)

• gemini_only (bool)

• document_task_type (str | None)

• query_task_type (str | None)

MAX_BATCH_SIZE = 50

MAX_BATCH_CHARS = 50000

MAX_EMBED_WORKERS = 20

__init__(api_key=None, model='google/gemini-embedding-001', dimensions=None, timeout=30.0, gemini_api_key=None, gemini_only=True, document_task_type=None, query_task_type=None)

Initialize the instance.

Parameters:

• api_key (Optional[str]) – Unused; kept for backward compatibility.

• model (str) – The model value.

• dimensions (Optional[int]) – The dimensions value.

• timeout (float) – Maximum wait time in seconds.

• gemini_api_key (Optional[str]) – Unused; pool is used instead.

• gemini_only (bool) – Unused; always Gemini API.

• document_task_type (Optional[str]) – Optional Gemini taskType for corpus (e.g. RETRIEVAL_DOCUMENT); used by embed_documents.

• query_task_type (Optional[str]) – Optional Gemini taskType for queries (e.g. RETRIEVAL_QUERY); used by embed_query.

name()

Name.

Returns:

Result string.

Return type:

str

dimension()

Dimension.

Returns:

The result.

Return type:

int

__call__(input)

ChromaDB EmbeddingFunction interface (legacy).

Uses document_task_type when set (same as embed_documents()).

Return type:

List[List[float]]

Parameters:

input (List[str])

embed_documents(input)

Embed documents (ChromaDB interface for upsert).

Return type:

List[List[float]]

Parameters:

input (List[str])

embed_query(input)

Embed query texts (ChromaDB interface for query).

Return type:

List[List[float]]

Parameters:

input (List[str])

class rag_system.FileRAGManager(store_name='default', store_path=None, api_key=None, embedding_model='google/gemini-embedding-001', max_file_size=15728640, gemini_only=True, document_task_type=None, query_task_type=None)

Bases: object

File-based RAG with ChromaDB storage and OpenRouter embeddings.

Parameters:

• store_name (str)

• store_path (str | None)

• api_key (str | None)

• embedding_model (str)

• max_file_size (int)

• gemini_only (bool)

• document_task_type (str | None)

• query_task_type (str | None)

__init__(store_name='default', store_path=None, api_key=None, embedding_model='google/gemini-embedding-001', max_file_size=15728640, gemini_only=True, document_task_type=None, query_task_type=None)

Initialize the instance.

Parameters:

• store_name (str) – The store name value.

• store_path (Optional[str]) – The store path value.

• api_key (Optional[str]) – The api key value.

• embedding_model (str) – The embedding model value.

• max_file_size (int) – The max file size value.

• gemini_only (bool) – Use only the Gemini API for embeddings.

• document_task_type (Optional[str]) – Optional Gemini task type for indexed text (e.g. RETRIEVAL_DOCUMENT).

• query_task_type (Optional[str]) – Optional Gemini task type for search queries (e.g. RETRIEVAL_QUERY).

index_file(file_path, tags=None, use_chunking=True, chunk_size=1500, chunk_overlap=200, force=False)

Index a single file into the collection.

When force is True the content-hash dedup check is skipped so the file is always re-embedded (but the store is not cleared).

Return type:

Dict[str, Any]

Parameters:

• file_path (str)

• tags (List[str] | None)

• use_chunking (bool)

• chunk_size (int)

• chunk_overlap (int)

• force (bool)

async index_url(url, tags=None, use_chunking=True, chunk_size=1500, chunk_overlap=200)

Index url.

Parameters:

• url (str) – URL string.

• tags (Optional[List[str]]) – The tags value.

• use_chunking (bool) – The use chunking value.

• chunk_size (int) – The chunk size value.

• chunk_overlap (int) – The chunk overlap value.

Returns:

The result.

Return type:

Dict[str, Any]

index_directory(directory_path, recursive=True, tags=None, exclude_patterns=None, max_workers=6, force=False, allowed_extensions=None)

Index all supported files in directory_path.

When max_workers > 1, files are indexed concurrently using a thread pool. Each file’s embedding batches are already parallelised inside the embedding function, so even max_workers=1 benefits from concurrent API calls.

force bypasses the per-file content-hash dedup check without clearing the store, so already-indexed files get re-embedded.

When allowed_extensions is set, only files whose suffix (after normalizing to a leading dot, lowercase) appears in the collection are queued; None means no extension filter (all supported types under SUPPORTED_EXTENSIONS).

Return type:

Dict[str, Any]

Parameters:

• directory_path (str)

• recursive (bool)

• tags (List[str] | None)

• exclude_patterns (List[str] | None)

• max_workers (int)

• force (bool)

• allowed_extensions (Collection[str] | None)

search(query, n_results=5, tags=None, return_content=True, query_embedding=None, max_content_size=8000)

Semantic search returning relevant chunks per file.

Instead of returning entire file contents, this collects the matching chunk texts that ChromaDB found and merges them (respecting max_content_size). Small files whose full text fits within one chunk are returned in full automatically.

Parameters:

• query (str) – Natural-language search query.

• n_results (int) – Maximum number of files to return.

• tags (Optional[List[str]]) – Optional tag filter.

• return_content (bool) – Include chunk text in results.

• query_embedding (list[float] | None) – Pre-computed query embedding (skips ChromaDB’s internal embedding call).

• max_content_size (int) – Maximum characters of merged chunk text to return per file (default 8000).

Return type:

List[Dict[str, Any]]

remove_file(file_path)

Delete the specified file.

Parameters:

file_path (str) – The file path value.

Returns:

The result.

Return type:

Dict[str, Any]

remove_url(url)

Delete the specified url.

Parameters:

url (str) – URL string.

Returns:

The result.

Return type:

Dict[str, Any]

list_indexed_files(limit=100)

List indexed files.

Parameters:

limit (int) – Maximum number of items.

Returns:

The result.

Return type:

List[Dict[str, Any]]

list_store_files()

List store files.

Returns:

The result.

Return type:

List[Dict[str, Any]]

read_store_file(filename)

Read store file.

Parameters:

filename (str) – The filename value.

Returns:

The result.

Return type:

Dict[str, Any]

close()

Release the ChromaDB client and its underlying SQLite resources.

Return type:

None

get_stats()

Retrieve the stats.

Returns:

The result.

Return type:

Dict[str, Any]

clear()

Clear.

Returns:

The result.

Return type:

Dict[str, Any]

rag_system.get_rag_store(store_name='default', api_key=None, max_file_size=None, gemini_only=True, document_task_type=None, query_task_type=None)

Get or create a RAG store by name (LRU-cached).

At most _STORE_REGISTRY_MAX_SIZE stores are kept open simultaneously. When a new store would exceed the limit the least recently used entry is closed and evicted.

Cache entries are keyed by store_name plus optional embedding task types so different embedding configurations do not share one client.

Return type:

FileRAGManager

Parameters:

• store_name (str)

• api_key (str | None)

• max_file_size (int | None)

• gemini_only (bool)

• document_task_type (str | None)

• query_task_type (str | None)

rag_system.get_stargazer_docs_store()

Return the shared RAG store for Sphinx / tool documentation.

Uses RETRIEVAL_DOCUMENT for indexed chunks and RETRIEVAL_QUERY for search queries (Gemini embedding task types).

Return type:

FileRAGManager

rag_system.list_rag_stores()

List all available RAG store directory names.

Return type:

List[str]

rag_system.list_rag_stores_with_stats()

List stores with file counts using only filesystem ops (no ChromaDB).

Counts physical files in each store’s files/ subdirectory as a lightweight proxy for the indexed entry count. This never opens a ChromaDB client and therefore uses zero additional RAM.

Return type:

List[Dict[str, Any]]

rag_system.delete_rag_store(store_name)

Delete a RAG store completely.

Return type:

Dict[str, Any]

Parameters:

store_name (str)

class rag_system.RAGAutoSearchManager(redis_client)

Bases: object

Per-channel auto-search configuration backed by Redis.

Parameters:

redis_client (aioredis.Redis)

__init__(redis_client)

Initialize the instance.

Parameters:

redis_client (Redis) – Redis connection client.

Return type:

None

async set_channel_config(channel_key, store_names, enabled=True, n_results=3, min_score=0.5)

Set auto-search configuration for a channel.

Parameters:

• channel_key (str) – Composite key "platform:channel_id".

• store_names (List[str])

• enabled (bool)

• n_results (int)

• min_score (float)

Return type:

Dict[str, Any]

async get_channel_config(channel_key)

Retrieve the channel config.

Parameters:

channel_key (str) – The channel key value.

Returns:

The result.

Return type:

Optional[Dict[str, Any]]

async disable_channel(channel_key)

Disable channel.

Parameters:

channel_key (str) – The channel key value.

Returns:

True on success, False otherwise.

Return type:

bool

async remove_channel_config(channel_key)

Delete the specified channel config.

Parameters:

channel_key (str) – The channel key value.

Returns:

True on success, False otherwise.

Return type:

bool

async list_configured_channels()

List configured channels.

Returns:

The result.

Return type:

List[Dict[str, Any]]

async search_for_message(channel_key, message_content, chunk_size=10000, query_embedding=None, user_id='')

Perform auto-search if the channel is configured.

Parameters:

• query_embedding (list[float] | None) – Pre-computed embedding for message_content. When provided it is forwarded to ChromaDB to skip its internal embedding call.

• user_id (str) – The message author. Used to enforce access control on cloud_usr_ stores.

• string (Returns XML-formatted RAG context)

• None. (or)

• channel_key (str)

• message_content (str)

• chunk_size (int)

Return type:

Optional[str]

Submodules

• rag_system.auto_search module
  • RAGAutoSearchManager
    • RAGAutoSearchManager.__init__()
    • RAGAutoSearchManager.set_channel_config()
    • RAGAutoSearchManager.get_channel_config()
    • RAGAutoSearchManager.disable_channel()
    • RAGAutoSearchManager.remove_channel_config()
    • RAGAutoSearchManager.list_configured_channels()
    • RAGAutoSearchManager.search_for_message()
• rag_system.file_rag_manager module
  • extract_pdf_text()
  • compress_pdf()
  • chunk_text()
  • fetch_url_content()
  • FileRAGManager
    • FileRAGManager.__init__()
    • FileRAGManager.index_file()
    • FileRAGManager.index_url()
    • FileRAGManager.index_directory()
    • FileRAGManager.search()
    • FileRAGManager.remove_file()
    • FileRAGManager.remove_url()
    • FileRAGManager.list_indexed_files()
    • FileRAGManager.list_store_files()
    • FileRAGManager.read_store_file()
    • FileRAGManager.close()
    • FileRAGManager.get_stats()
    • FileRAGManager.clear()
  • get_rag_store()
  • get_stargazer_docs_store()
  • list_rag_stores()
  • list_rag_stores_with_stats()
  • delete_rag_store()
• rag_system.openrouter_embeddings module
  • OpenRouterEmbeddings
    • OpenRouterEmbeddings.DEFAULT_MODEL
    • OpenRouterEmbeddings.MAX_BATCH_SIZE
    • OpenRouterEmbeddings.MAX_BATCH_CHARS
    • OpenRouterEmbeddings.__init__()
    • OpenRouterEmbeddings.embed_text()
    • OpenRouterEmbeddings.embed_texts()
    • OpenRouterEmbeddings.embed_text_for_search()
    • OpenRouterEmbeddings.close()
    • OpenRouterEmbeddings.__aenter__()
    • OpenRouterEmbeddings.__aexit__()
  • SyncOpenRouterEmbeddings
    • SyncOpenRouterEmbeddings.MAX_BATCH_SIZE
    • SyncOpenRouterEmbeddings.MAX_BATCH_CHARS
    • SyncOpenRouterEmbeddings.MAX_EMBED_WORKERS
    • SyncOpenRouterEmbeddings.__init__()
    • SyncOpenRouterEmbeddings.name()
    • SyncOpenRouterEmbeddings.dimension()
    • SyncOpenRouterEmbeddings.__call__()
    • SyncOpenRouterEmbeddings.embed_documents()
    • SyncOpenRouterEmbeddings.embed_query()