railtracks.prebuilt.rag_node 

 1from typing import List
 2
 3import railtracks as rt
 4from railtracks.rag.rag_core import RAG, RAGConfig, SearchResult
 5
 6
 7def rag_node(
 8    documents: List[str],
 9    embed_model="text-embedding-3-small",
10    token_count_model="gpt-4o",
11    chunk_size=1000,
12    chunk_overlap=200,
13):
14    """
15    Creates a rag node that allows you to vector the search the provided documents.
16
17    Args:
18        documents (list): List of documents to process. Each document is a string of raw text.
19        embed_model (str): Model name for embedding service.
20        token_count_model (str): Model name for token counting.
21        chunk_size (int): Size of each text chunk.
22        chunk_overlap (int): Overlap between chunks.
23
24    Returns:
25        A node to be invoked upon request.
26
27    """
28
29    rag_core = RAG(
30        docs=documents,
31        config=RAGConfig(
32            embedding={"model": embed_model},
33            store={},
34            chunking={
35                "chunk_size": chunk_size,
36                "chunk_overlap": chunk_overlap,
37                "model": token_count_model,
38            },
39        ),
40    )
41    rag_core.embed_all()
42
43    def query(query: str, top_k: int = 1) -> SearchResult:
44        result = rag_core.search(query, top_k=top_k)
45        return result
46
47    return rt.function_node(query)
def rag_node( documents: List[str], embed_model='text-embedding-3-small', token_count_model='gpt-4o', chunk_size=1000, chunk_overlap=200): 
 8def rag_node(
 9    documents: List[str],
10    embed_model="text-embedding-3-small",
11    token_count_model="gpt-4o",
12    chunk_size=1000,
13    chunk_overlap=200,
14):
15    """
16    Creates a rag node that allows you to vector the search the provided documents.
17
18    Args:
19        documents (list): List of documents to process. Each document is a string of raw text.
20        embed_model (str): Model name for embedding service.
21        token_count_model (str): Model name for token counting.
22        chunk_size (int): Size of each text chunk.
23        chunk_overlap (int): Overlap between chunks.
24
25    Returns:
26        A node to be invoked upon request.
27
28    """
29
30    rag_core = RAG(
31        docs=documents,
32        config=RAGConfig(
33            embedding={"model": embed_model},
34            store={},
35            chunking={
36                "chunk_size": chunk_size,
37                "chunk_overlap": chunk_overlap,
38                "model": token_count_model,
39            },
40        ),
41    )
42    rag_core.embed_all()
43
44    def query(query: str, top_k: int = 1) -> SearchResult:
45        result = rag_core.search(query, top_k=top_k)
46        return result
47
48    return rt.function_node(query)

Creates a rag node that allows you to vector the search the provided documents.

Arguments:
  • documents (list): List of documents to process. Each document is a string of raw text.
  • embed_model (str): Model name for embedding service.
  • token_count_model (str): Model name for token counting.
  • chunk_size (int): Size of each text chunk.
  • chunk_overlap (int): Overlap between chunks.
Returns:

A node to be invoked upon request.