ComoRAG Framework

To run ComoRAG locally with vLLM, follow these three steps:

1. Start the vLLM Server

You can start the server using either the vllm serve command or the python entrypoint module.

Option A: vllm serve

vllm serve /path/to/your/model \
  --tensor-parallel-size 1 \
  --max-model-len 4096 \
  --gpu-memory-utilization 0.95

Option B: python module

python -m vllm.entrypoints.openai.api_server \
  --model /path/to/your/model \
  --served-model-name your-model-name \
  --tensor-parallel-size 1 \
  --max-model-len 32768 \
  --dtype auto

2. Configure main_vllm.py

Update the BaseConfig in main_vllm.py to point to your local server address and model name.

vllm_base_url = 'http://localhost:8000/v1'
served_model_name = '/path/to/your/model'

config = BaseConfig(
    llm_base_url=vllm_base_url,
    llm_name=served_model_name,
    llm_api_key="your-api-key-here",
    dataset='cinderella',
    embedding_model_name='/path/to/your/embedding/model',
    embedding_batch_size=4,
    need_cluster=True,
    output_dir='result/cinderella_vllm',
    save_dir='outputs/cinderella_vllm',
    max_meta_loop_max_iterations=5,
    is_mc=False,
    max_tokens_ver=2000,
    max_tokens_sem=2000,
    max_tokens_epi=2000
)

3. Run the program

python main_vllm.py

To add new prompt templates to ComoRAG, create a new .py file within the src/comorag/prompts/templates/ directory. Each file (excluding __init__.py) is treated as a module under the narrtiverag project.

To make a template accessible, you must define a variable named prompt_template within the file. The filename (without the .py extension) will serve as the unique key used to access that template in the system.

Supported template formats for the prompt_template variable:

1. String: A standard string, optionally containing ${placeholder} syntax for value injection (these are automatically processed using Template(...)).
2. Template Instance: A pre-instantiated Template object.
3. Chat History: A list of dictionaries representing a conversation, where each dictionary follows the format: {"role": "system" | "user" | "assistant", "content": "..."}.

# Example: String template with placeholders
# Filename: my_custom_template.py

prompt_template = "You are a helpful assistant. Answer the following: ${query}"

# Example: Chat history template
# Filename: chat_template.py

prompt_template = [
    {"role": "system", "content": "You are a helpful assistant."},
    {"role": "user", "content": "Hello!"}
]

To use remote LLMs via OpenAI, use the main_openai.py script. You must configure a BaseConfig object within the script to specify your API endpoint, model names, and dataset paths.

Configuration Steps

1. Edit main_openai.py to set your llm_base_url, llm_name, dataset, and embedding_model_name.
2. Set max_meta_loop_max_iterations to control the reasoning loop depth.
3. Run the script.

Execution

python main_openai.py

config = BaseConfig(
    llm_base_url='https://api.example.com/v1',  # OpenAI API
    llm_name='gpt-4o-mini',
    dataset='cinderella',
    embedding_model_name='/path/to/your/embedding/model',
    embedding_batch_size=32,
    need_cluster=True,  # Enable semantic/plot enhancement
    output_dir='result/cinderella',
    save_dir='outputs/cinderella',
    max_meta_loop_max_iterations=5,  # Max iterations
    is_mc=False,  # False if not multiple choice
    max_tokens_ver=2000,
    max_tokens_sem=2000,
    max_tokens_epi=2000
)

The make_cache_embed utility provides a decorator-like wrapper to cache embedding results to a SQLite database. This prevents redundant calls to expensive embedding APIs or models.

How it works:

1. It generates a unique SHA256 hash for every prompt based on the instruction, prompts, and max_length provided in kwargs.
2. It checks a SQLite database (cache_file_name) for an existing hash.
3. If found, it returns the cached embedding.
4. If not found, it calls the provided encode_func, then saves the new embeddings to the database using a file lock to ensure concurrency safety.

Requirements for encode_func:

• The function must accept prompts (a list of strings) and return a 2D torch.Tensor or a NumPy array that can be converted to a tensor.
• It should also accept instruction and max_length in its kwargs to ensure the cache key is consistent.

from comorag.embedding_model.base import make_cache_embed
import torch

# An example encoder function
def my_encoder(prompts, instruction="", max_length=512, **kwargs):
    # Returns a torch.Tensor of shape (len(prompts), dim)
    return torch.randn(len(prompts), 768)

# Wrap the encoder with SQLite caching
cached_encode = make_cache_embed(
    encode_func=my_encoder,
    cache_file_name="my_embeddings.db",
    device="cpu"
)

# First call: computes and saves to DB
vecs = cached_encode(prompts=["query 1", "query 2"], instruction="embed")

# Second call: retrieves from DB
vecs_cached = cached_encode(prompts=["query 1", "query 2"], instruction="embed")

Memory fusion allows the system to synthesize related historical nodes into a single, coherent FUSION type node using an LLM agent.

The Fusion Workflow:

1. create_fusion_content(probe, top_k_percent):
   • Retrieves similar nodes from the pool based on the probe.
   • Formats the cue of those nodes into a structured string.
   • Calls agent.fuse_memory_nodes(query, content, max_completion_tokens) to generate the fused text.
2. add_fused_node(probe, fused_content, source_nodes):
   • Creates a new MemoryNode with node_type=NodeType.FUSION.
   • Sets the fused_content as the node's cue.
   • Computes the embedding for the new node.
   • Automatically adds the new node to the temporary pool (temp_pool). You must call merge_temp_to_main() later to persist it.

# 1. Generate fused text using the agent
fused_text = pool.create_fusion_content(probe="Summarize recent findings", top_k_percent=0.2)

# 2. Create the node and add it to the temp_pool
# Note: source_nodes is passed for tracking/context
source_nodes = pool.retrieve_similar_nodes("Summarize recent findings", 0.2)
pool.add_fused_node(probe="Summarize recent findings", fused_content=fused_text, source_nodes=source_nodes)

# 3. Persist the fusion to the main pool
pool.merge_temp_to_main()